├── .github
└── workflows
│ └── pythonpublish.yml
├── .gitignore
├── .pre-commit-config.yaml
├── CHANGELOG.md
├── LICENSE
├── README.md
├── docker-compose.yml
├── examples
├── producer.py
└── worker.py
├── faktory
├── __init__.py
├── _proto.py
├── client.py
├── exceptions.py
└── worker.py
├── setup.cfg
├── setup.py
└── tests
├── test_client.py
├── test_exceptions.py
├── test_formatting.py
├── test_proto.py
└── test_worker.py
/.github/workflows/pythonpublish.yml:
--------------------------------------------------------------------------------
1 | name: Upload Python Package
2 |
3 | on:
4 | release:
5 | types: [ published ]
6 |
7 | permissions:
8 | contents: read
9 |
10 | jobs:
11 | deploy:
12 | runs-on: ubuntu-latest
13 | steps:
14 | - uses: actions/checkout@v3
15 |
16 | - name: Set up Python
17 | uses: actions/setup-python@v3
18 | with:
19 | python-version: '3.x'
20 |
21 | - name: Install dependencies
22 | run: |
23 | python -m pip install --upgrade pip
24 | pip install build
25 |
26 | - name: Build package
27 | run: python -m build
28 |
29 | - name: Publish package
30 | uses: pypa/gh-action-pypi-publish@release/v1
31 | with:
32 | user: __token__
33 | password: ${{ secrets.PYPI_API_TOKEN }}
34 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | # Byte-compiled / optimized / DLL files
2 | __pycache__/
3 | *.py[cod]
4 | *$py.class
5 |
6 | /*.egg-info
7 | /build
8 | /dist
9 |
10 | # Editors
11 | .vscode/
12 | .idea/
13 |
14 | # Virtual Environments
15 | venv/
16 | env/
17 | .venv/
--------------------------------------------------------------------------------
/.pre-commit-config.yaml:
--------------------------------------------------------------------------------
1 | repos:
2 | - repo: local
3 | hooks:
4 | - id: black
5 | name: black
6 | entry: black
7 | language: python
8 | language_version: python3
9 | types: [python]
10 |
--------------------------------------------------------------------------------
/CHANGELOG.md:
--------------------------------------------------------------------------------
1 | # [Changelog](https://github.com/cdrx/faktory_worker_python/releases)
2 |
3 | All notable changes to this project will be documented in this file.
4 |
5 | ## 1.0.0 - 2022-04-30
6 |
7 | - Python 3.7+ is now required (#36)
8 | - Fixed a type issue that crashed the worker (#41)
9 | - Improved recovery from broken process pool error, thanks to @tswayne (#42)
10 |
11 | ## 0.5.1 - 2021-03-20
12 |
13 | - Fixed issue with installing 0.5.0
14 |
15 | ## 0.5.0 - 2021-03-01
16 |
17 | - This release was yanked due to a build issue -- use 0.5.1 instead
18 | - Added support for the `backtrace` option, thanks to @tswayne
19 |
20 | ## 0.4.0 - 2018-02-21
21 |
22 | - Added support for job priorities
23 | - Added validation if args isn't something faktory will accept
24 | - Added a pool of threads based worker implementation
25 | - Compatibility with Faktory 0.7
26 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | BSD 3-Clause License
2 |
3 | Copyright (c) 2017, Chris R
4 | All rights reserved.
5 |
6 | Redistribution and use in source and binary forms, with or without
7 | modification, are permitted provided that the following conditions are met:
8 |
9 | * Redistributions of source code must retain the above copyright notice, this
10 | list of conditions and the following disclaimer.
11 |
12 | * Redistributions in binary form must reproduce the above copyright notice,
13 | this list of conditions and the following disclaimer in the documentation
14 | and/or other materials provided with the distribution.
15 |
16 | * Neither the name of the copyright holder nor the names of its
17 | contributors may be used to endorse or promote products derived from
18 | this software without specific prior written permission.
19 |
20 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
21 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
22 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
23 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
24 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
25 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
26 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
27 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
28 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
30 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Python worker for Faktory
2 |
3 | ## Overview
4 |
5 | This project is a complete worker and client implementation for the [Faktory job server](https://github.com/contribsys/faktory). You can use it to either consume jobs from Faktory or push jobs to the Faktory server to be processed.
6 |
7 | Requires Python 3.7+.
8 |
9 | #### Supported Faktory Versions
10 |
11 | :x: 0.5.0
12 | :white_check_mark: 0.6
13 | :white_check_mark: 0.7
14 | :white_check_mark: 0.8
15 | :white_check_mark: 1.0 and up
16 |
17 | ## Features
18 |
19 | - [x] Creating a worker to run jobs from Faktory
20 | - [x] Concurrency (with multiple processes or threads with the `use_threads=True` option)
21 | - [x] Pushing work to Faktory from Python (with retries, custom metadata and scheduled support)
22 | - [x] Pushing exception / errors from Python back up to Faktory
23 | - [x] Sends worker status back to Faktory
24 | - [x] Supports quiet and teminate from the Faktory web UI
25 | - [x] Password authentication
26 | - [x] TLS support
27 | - [x] Graceful worker shutdown (ctrl-c will allow 15s for pending jobs to finish)
28 |
29 | #### Todo
30 |
31 | - [ ] Documentation (in progress, help would be appreciated)
32 | - [ ] Tests (in progress, help would be appreciated)
33 | - [ ] Django integration (`./manage.py runworker` and `app/tasks.py` support)
34 |
35 | ## Installation
36 |
37 | ```
38 | pip install faktory
39 | ```
40 |
41 | ## Pushing Work to Faktory
42 |
43 | There is a client context manager that you can use like this:
44 |
45 | ```
46 | import faktory
47 |
48 | with faktory.connection() as client:
49 | client.queue('test', args=(1, 2))
50 | client.queue('test', args=(4, 5), queue='other')
51 | ```
52 |
53 | `test` doesn't need to be implemented by the Python worker, it can be any of the available worker implementations.
54 |
55 | ## Worker Example
56 |
57 | To create a faktory worker (to process jobs from the server) you'll need something like this:
58 |
59 | ```
60 | from faktory import Worker
61 |
62 | def your_function(x, y):
63 | return x + y
64 |
65 | w = Worker(queues=['default'], concurrency=1)
66 | w.register('test', your_function)
67 |
68 | w.run() # runs until control-c or worker shutdown from Faktory web UI
69 |
70 | ```
71 | ## Concurrency
72 |
73 | The default mode of concurrency is to use a [ProcessPoolExecutor](https://devdocs.io/python~3.11/library/concurrent.futures#concurrent.futures.ProcessPoolExecutor). Multiple processes are started, the number being controlled by the `concurrency` keyword argument of the `Worker` class. New processes are started only once, and stay up, processing jobs from the queue. There is the possibility to use threads instead of processes as a concurency mechanism. This is done by using `use_threads=True` at Worker creation. As with processes, threads are started once and reused for each job. When doing so, be mindful of the consequences of using threads in your code, like global variables concurrent access, or the fact that initialization code that is run outside of the registered functions will be run only once at worker startup, not once for each thread.
74 |
75 | #### Samples
76 |
77 | There is very basic [example worker](examples/worker.py) and an [example producer](examples/producer.py) that you can use as a basis for your project.
78 |
79 | #### Connection to Faktory
80 |
81 | faktory_worker_python uses this format for the Faktory URL:
82 |
83 | `tcp://:password@localhost:7419`
84 |
85 | or with TLS:
86 |
87 | `tcp+tls://:password@localhost:7419`
88 |
89 | If the environment variable `FAKTORY_URL` is set, that is used. Otherwise you can pass the server URL in to the `Worker` or `Client` constructor, like this:
90 |
91 | ```w = Worker(faktory="tcp://localhost:7419")```
92 |
93 | #### Logging
94 |
95 | The worker users Python's built in logging module, which you can enable like this before calling `.run()`:
96 |
97 | ```
98 | import logging
99 | logging.basicConfig(level=logging.DEBUG)
100 | ```
101 |
102 | ## Troubleshooting
103 |
104 | ### Registering decorated functions
105 |
106 | When using the default multiprocessing mode of concurrency, the underlying process pool uses the standard library's `pickle` module to serialize registered functions. However, a function can only be `pickled` if defined directly at the top-level of a module. If the function is instead produced by a decorator, the pickling won't work. A workaround for this issue is to change the mode of concurrenry and use threads instead:
107 |
108 | ```
109 | w = Worker(..., use_threads=True)
110 | ```
111 |
--------------------------------------------------------------------------------
/docker-compose.yml:
--------------------------------------------------------------------------------
1 | version: "3"
2 |
3 | services:
4 | faktory:
5 | image: contribsys/faktory:latest
6 | ports:
7 | - "7419:7419"
8 | - "7420:7420"
9 |
--------------------------------------------------------------------------------
/examples/producer.py:
--------------------------------------------------------------------------------
1 | import time
2 |
3 | import faktory
4 |
5 | with faktory.connection() as client:
6 | while True:
7 | client.queue("add", args=(1, 2), queue="default")
8 | time.sleep(1)
9 |
10 | client.queue("subtract", args=(10, 5), queue="default")
11 | time.sleep(1)
12 |
13 | client.queue("multiply", args=(8, 8), queue="default")
14 | time.sleep(1)
15 |
--------------------------------------------------------------------------------
/examples/worker.py:
--------------------------------------------------------------------------------
1 | import logging
2 |
3 | logging.basicConfig(level=logging.INFO)
4 |
5 | from faktory import Worker
6 |
7 |
8 | def add_numbers(x, y):
9 | calc = x + y
10 |
11 | print(f"add: {x} + {y} = {calc}")
12 | return calc
13 |
14 |
15 | def subtract_numbers(x, y):
16 | calc = x - y
17 |
18 | print(f"subtract: {x} - {y} = {calc}")
19 | return calc
20 |
21 |
22 | def multiply_numbers(x, y):
23 | calc = x * y
24 |
25 | print(f"multiply: {x} * {y} = {calc}")
26 | return calc
27 |
28 |
29 | if __name__ == "__main__":
30 | w = Worker(faktory="tcp://localhost:7419", queues=["default"], concurrency=1)
31 | w.register("add", add_numbers)
32 | w.register("subtract", subtract_numbers)
33 | w.register("multiply", multiply_numbers)
34 | w.run() # runs until control-c or worker shutdown from Faktory web UI
35 |
36 | # check examples/producer.py for how to submit tasks to be run by faktory
37 |
--------------------------------------------------------------------------------
/faktory/__init__.py:
--------------------------------------------------------------------------------
1 | from contextlib import contextmanager
2 |
3 | from .client import Client
4 | from .exceptions import *
5 | from .worker import Worker
6 |
7 | __version__ = "1.0.0"
8 | __url__ = "https://github.com/cdrx/faktory_worker_python"
9 |
10 |
11 | def get_client(*args, **kwargs):
12 | return Client(*args, **kwargs)
13 |
14 |
15 | @contextmanager
16 | def connection(*args, **kwargs):
17 | c = get_client(*args, **kwargs)
18 | c.connect()
19 | yield c
20 | c.disconnect()
21 |
--------------------------------------------------------------------------------
/faktory/_proto.py:
--------------------------------------------------------------------------------
1 | import hashlib
2 | import json
3 | import logging
4 | import os
5 | import os.path
6 | import select
7 | import socket
8 | import ssl
9 | from typing import Any, Dict, Iterable, Iterator, List, Optional
10 | from urllib.parse import urlparse
11 |
12 | from .exceptions import (
13 | FaktoryAuthenticationError,
14 | FaktoryConnectionResetError,
15 | FaktoryHandshakeError,
16 | )
17 |
18 |
19 | class Connection:
20 | buffer_size = 4096
21 | timeout = 30
22 | use_tls = False
23 | send_heartbeat_every = 15
24 | labels = ["python"]
25 | queues = ["default"]
26 | debug = False
27 |
28 | is_connected = False
29 | is_connecting = False
30 | is_quiet = False
31 | is_disconnecting = False
32 | disconnection_requested = None
33 | force_disconnection_after = None
34 |
35 | def __init__(
36 | self,
37 | faktory: Optional[str] = None,
38 | timeout: int = 30,
39 | buffer_size: int = 4096,
40 | worker_id=None,
41 | labels: List[str] = None,
42 | log: logging.Logger = None,
43 | ):
44 | if not faktory:
45 | faktory = os.environ.get("FAKTORY_URL", "tcp://localhost:7419")
46 |
47 | url = urlparse(faktory)
48 | self.host = url.hostname
49 | self.port = url.port or 7419
50 | self.password = url.password
51 |
52 | if "tls" in url.scheme:
53 | self.use_tls = True
54 |
55 | self.timeout = timeout
56 | self.buffer_size = buffer_size
57 |
58 | self.labels = labels
59 | if not self.labels:
60 | self.labels = []
61 |
62 | self.worker_id = worker_id
63 | self.socket = None
64 |
65 | self.log = log or logging.getLogger(name="faktory.connection")
66 |
67 | def connect(self, worker_id=None) -> bool:
68 | self.log.info("Connecting to {}:{}".format(self.host, self.port))
69 | self.is_connecting = True
70 |
71 | self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
72 | if self.use_tls:
73 | self.log.debug("Using TLS")
74 | self.socket = ssl.wrap_socket(self.socket)
75 |
76 | self.socket.setblocking(0)
77 | self.socket.settimeout(self.timeout)
78 | try:
79 | self.socket.connect((self.host, self.port))
80 | except ssl.SSLError:
81 | raise
82 |
83 | ahoy = next(self.get_message())
84 | if not ahoy.startswith("HI "):
85 | raise FaktoryHandshakeError(
86 | "Could not connect to Faktory; expected HI from server, but got '{}'".format(
87 | ahoy
88 | )
89 | )
90 |
91 | response = {
92 | "hostname": socket.gethostname(),
93 | "pid": os.getpid(),
94 | "labels": self.labels,
95 | }
96 |
97 | if worker_id:
98 | response["wid"] = self.worker_id
99 |
100 | try:
101 | handshake = json.loads(ahoy[len("HI ") :])
102 | version = int(handshake["v"])
103 | if not self.is_supported_server_version(version):
104 | self.socket.close()
105 | raise FaktoryHandshakeError(
106 | "Could not connect to Faktory; unsupported server version {}".format(
107 | version
108 | )
109 | )
110 |
111 | nonce = handshake.get("s")
112 | if nonce and self.password:
113 | response["pwdhash"] = hashlib.sha256(
114 | str.encode(self.password) + str.encode(nonce)
115 | ).hexdigest()
116 | except (ValueError, TypeError):
117 | self.socket.close()
118 | raise FaktoryHandshakeError(
119 | "Could not connect to Faktory; expected handshake format"
120 | )
121 |
122 | self.reply("HELLO", response)
123 |
124 | ok = next(self.get_message())
125 | if ok != "OK":
126 | if ok.startswith("ERR") and "invalid password" in ok.lower():
127 | self.socket.close()
128 | raise FaktoryAuthenticationError(
129 | "Could not connect to Faktory; wrong password"
130 | )
131 | self.socket.close()
132 | raise FaktoryHandshakeError(
133 | "Could not connect to Faktory; expected OK from server, but got '{}'".format(
134 | ok
135 | )
136 | )
137 |
138 | self.log.debug("Connected to Faktory")
139 |
140 | self.is_connected, self.is_connecting = True, False
141 | return self.is_connected
142 |
143 | def _validate_handshake(self, payload: Dict[str, Any]) -> None:
144 | """
145 | Helper function designed to authenticate with the Faktory
146 | server and validate the connection is as expected.
147 |
148 | Args:
149 | - payload (Dict[str, Any]): Response from sending the initial `Hello` message. Payload
150 | is expected to have the following keys: `[hostname, pid, labels]` and
151 | optionally have the following keys: `[worker_id, pwdhash]`
152 |
153 | Raises:
154 | - FaktoryHandshakeError: Raised when receiving an unexpected response
155 | in the handshake message from the server.
156 | - FaktoryAuthenticationError: Raised when an invalid password is supplied.
157 |
158 | Response:
159 | - None
160 | """
161 | self.reply("HELLO", payload)
162 |
163 | ok = next(self.get_message())
164 | if ok != "OK":
165 | if ok.startswith("ERR") and "invalid password" in ok.lower():
166 | self.socket.close()
167 | raise FaktoryAuthenticationError(
168 | "Could not connect to Faktory; wrong password"
169 | )
170 | self.socket.close()
171 | raise FaktoryHandshakeError(
172 | "Could not connect to Faktory; expected OK from server, but got '{}'".format(
173 | ok
174 | )
175 | )
176 |
177 | self.log.debug("Connected to Faktory")
178 |
179 | def validate_connection(self, send_worker_id: bool = False) -> None:
180 | """
181 | Connects to the Faktory job server and validates the connection.
182 |
183 | Args:
184 | - send_worker_id (bool): Flag of whether to send the worker_id to the
185 | Faktory server.
186 |
187 | Raises:
188 | - FaktoryHandshakeError: Raised if a message with an unexpected
189 | format is received from the job server.
190 | - FaktoryHandshakeError: Raised if the job server's version
191 | isn't supported by this client.
192 | - FaktoryAuthenticationError: Raised if invalid credentials are
193 | used when trying to authenticate with the job server.
194 |
195 | Returns:
196 | - None
197 | """
198 | ahoy = next(self.get_message())
199 | if not ahoy.startswith("HI "):
200 | raise FaktoryHandshakeError(
201 | "Could not connect to Faktory; expected HI from server, but got '{}'".format(
202 | ahoy
203 | )
204 | )
205 |
206 | response = {
207 | "hostname": socket.gethostname(),
208 | "pid": os.getpid(),
209 | "labels": self.labels,
210 | }
211 |
212 | if send_worker_id:
213 | response["wid"] = self.worker_id
214 |
215 | try:
216 | handshake = json.loads(ahoy[len("HI ") :])
217 | version = int(handshake["v"])
218 | if not self.is_supported_server_version(version):
219 | self.socket.close()
220 | raise FaktoryHandshakeError(
221 | "Could not connect to Faktory; unsupported server version {}".format(
222 | version
223 | )
224 | )
225 |
226 | nonce = handshake.get("s")
227 | if nonce and self.password:
228 | response["pwdhash"] = hashlib.sha256(
229 | str.encode(self.password) + str.encode(nonce)
230 | ).hexdigest()
231 | except (ValueError, TypeError):
232 | self.socket.close()
233 | raise FaktoryHandshakeError(
234 | "Could not connect to Faktory; expected handshake format"
235 | )
236 | self._validate_handshake(response)
237 |
238 | self.is_connected, self.is_connecting = True, False
239 | return self.is_connected
240 |
241 | def is_supported_server_version(self, v: int) -> bool:
242 | return v == 2
243 |
244 | def get_message(self) -> Iterator[str]:
245 | socket = self.socket
246 | buffer = self.select_data(self.buffer_size)
247 | while self.is_connected or self.is_connecting:
248 | buffering = True
249 | while buffering:
250 | if buffer.count(b"\r\n"):
251 | (line, buffer) = buffer.split(b"\r\n", 1)
252 | if len(line) == 0:
253 | continue
254 | elif chr(line[0]) == "+":
255 | resp = line[1:].decode().strip("\r\n ")
256 | if self.debug:
257 | self.log.debug("> {}".format(resp))
258 | yield resp
259 | elif chr(line[0]) == "-":
260 | resp = line[1:].decode().strip("\r\n ")
261 | if self.debug:
262 | self.log.debug("> {}".format(resp))
263 | yield resp
264 | elif chr(line[0]) == "$":
265 | # read $xxx bytes of data into a buffer
266 | number_of_bytes = (
267 | int(line[1:]) + 2
268 | ) # add 2 bytes so we read the \r\n from the end
269 | if number_of_bytes <= 1:
270 | if self.debug:
271 | self.log.debug("> {}".format("nil"))
272 | yield ""
273 | else:
274 | if len(buffer) >= number_of_bytes:
275 | # we've already got enough bytes in the buffer
276 | data = buffer[:number_of_bytes]
277 | buffer = buffer[number_of_bytes:]
278 | else:
279 | data = buffer
280 | while len(data) != number_of_bytes:
281 | bytes_required = number_of_bytes - len(data)
282 | data += self.select_data(bytes_required)
283 | buffer = []
284 | resp = data.decode().strip("\r\n ")
285 | if self.debug:
286 | self.log.debug("> {}".format(resp))
287 | yield resp
288 | else:
289 | more = self.select_data(self.buffer_size)
290 | if not more:
291 | buffering = False
292 | else:
293 | buffer += more
294 |
295 | def select_data(self, buffer_size: int):
296 | s = self.socket
297 | ready = select.select([s], [], [], self.timeout)
298 | if ready[0]:
299 | buffer = s.recv(buffer_size)
300 | if len(buffer) > 0:
301 | return buffer
302 | self.disconnect()
303 | raise FaktoryConnectionResetError
304 |
305 | def fetch(self, queues: Iterable[str]) -> Optional[dict]:
306 | self.reply("FETCH {}".format(" ".join(queues)))
307 | job = next(self.get_message())
308 | if not job:
309 | return None
310 |
311 | data = json.loads(job)
312 | return data
313 |
314 | def reply(self, cmd: str, data: Any = None):
315 | if self.debug:
316 | self.log.debug("< {} {}".format(cmd, data or ""))
317 | s = cmd
318 | if data is not None:
319 | if type(data) is dict:
320 | s = "{} {}".format(s, json.dumps(data))
321 | else:
322 | s = "{} {}".format(s, data)
323 | buffer = str.encode(s + "\r\n")
324 | while len(buffer):
325 | sent = self.socket.send(buffer)
326 | if sent == 0:
327 | raise FaktoryConnectionResetError
328 | buffer = buffer[sent:]
329 |
330 | def disconnect(self):
331 | self.log.info("Disconnected")
332 | self.socket.close()
333 | self.is_connected = False
334 |
--------------------------------------------------------------------------------
/faktory/client.py:
--------------------------------------------------------------------------------
1 | import typing
2 | import uuid
3 |
4 | from ._proto import Connection
5 |
6 |
7 | class Client:
8 | is_connected = False
9 |
10 | def __init__(self, faktory=None, connection=None, **kwargs):
11 | self.faktory = connection or Connection(faktory, **kwargs)
12 |
13 | def __enter__(self):
14 | self.connect()
15 | return self
16 |
17 | def __exit__(self, *args):
18 | self.disconnect()
19 |
20 | def connect(self):
21 | self.is_connected = self.faktory.connect()
22 | return self.is_connected
23 |
24 | def disconnect(self):
25 | self.faktory.disconnect()
26 | self.is_connected = False
27 |
28 | def queue(
29 | self,
30 | task: str,
31 | args: typing.Iterable = None,
32 | queue: str = "default",
33 | priority: int = 5,
34 | jid: str = None,
35 | custom=None,
36 | reserve_for=None,
37 | at=None,
38 | retry=5,
39 | backtrace=0,
40 | ):
41 | was_connected = self.is_connected
42 | if not self.is_connected:
43 | # connect if we are not already connected
44 | self.connect()
45 |
46 | if not task:
47 | raise ValueError("Empty task name")
48 |
49 | if not queue:
50 | raise ValueError("Empty queue name")
51 |
52 | if not jid:
53 | jid = self.random_job_id()
54 |
55 | if args is None:
56 | args = ()
57 |
58 | request = {"jid": jid, "queue": queue, "jobtype": task, "priority": priority}
59 |
60 | if custom is not None:
61 | request["custom"] = custom
62 |
63 | if args is not None:
64 | if not isinstance(
65 | args, (typing.Iterator, typing.Set, typing.List, typing.Tuple)
66 | ):
67 | raise ValueError(
68 | "Argument `args` must be an iterator, generator, list, tuple or a set"
69 | )
70 |
71 | request["args"] = list(args)
72 |
73 | if reserve_for is not None:
74 | request["reserve_for"] = reserve_for
75 |
76 | if at is not None:
77 | request["at"] = at
78 |
79 | request["retry"] = retry
80 | request["backtrace"] = backtrace
81 |
82 | self.faktory.reply("PUSH", request)
83 | ok = next(self.faktory.get_message())
84 |
85 | if not was_connected:
86 | self.disconnect()
87 |
88 | return ok == "OK"
89 |
90 | def random_job_id(self) -> str:
91 | return uuid.uuid4().hex
92 |
--------------------------------------------------------------------------------
/faktory/exceptions.py:
--------------------------------------------------------------------------------
1 | __ALL__ = ["FaktoryHandshakeError", "FaktoryAuthenticationError", "FaktoryError"]
2 |
3 |
4 | class FaktoryError(Exception):
5 | """
6 | The base Faktory Exception that all other Faktory related
7 | exceptions inherit from. To catch any other, more specific,
8 | Faktory exception, catch this one and handle appropriately.
9 |
10 | This Exception should not be used, and is in place strictly
11 | as a way to catch all `Faktory` exceptions.
12 | """
13 |
14 | pass
15 |
16 |
17 | class FaktoryHandshakeError(FaktoryError, ConnectionError):
18 | pass
19 |
20 |
21 | class FaktoryAuthenticationError(FaktoryError, ConnectionError):
22 | pass
23 |
24 |
25 | class FaktoryConnectionResetError(FaktoryError, ConnectionResetError):
26 | pass
27 |
--------------------------------------------------------------------------------
/faktory/worker.py:
--------------------------------------------------------------------------------
1 | import logging
2 | import signal
3 | import sys
4 | import time
5 | import uuid
6 | from collections import namedtuple
7 | from concurrent.futures import Executor, ProcessPoolExecutor, ThreadPoolExecutor, BrokenExecutor
8 | from concurrent.futures.process import BrokenProcessPool
9 | from concurrent.futures.thread import BrokenThreadPool
10 | from datetime import datetime, timedelta
11 | from typing import Callable, Iterable
12 |
13 | from ._proto import Connection
14 |
15 | Task = namedtuple("Task", ["name", "func", "bind"])
16 |
17 |
18 | class Worker:
19 | send_heartbeat_every = 15 # seconds
20 | is_quiet = False
21 | is_disconnecting = False
22 |
23 | def __init__(self, *args, **kwargs):
24 | """
25 | Creates a Faktory worker.
26 |
27 | This worker will connect to the `faktory` argument by default. It should be in the standard Faktory format:
28 | ```
29 | tcp://:password@localhost:7419
30 | ```
31 | If you don't pass a faktory instance to connect to, the worker will check the `FAKTORY_URL` environment variable.
32 | If the environment variable is not set, then the worker will attempt to connect to Faktory on the localhost,
33 | without a password.
34 |
35 | If the URL scheme is `tcp+tls://` then the Faktory worker will establish a TLS encrypted connection to Faktory.
36 |
37 | You may pass a list of queues to process with the `queues` argument. If you supply no `queues`, then the worker
38 | will process the default queue.
39 |
40 | You may pass a list of labels to process with the `labels` argument. These are visible in the Faktory Web UI. If
41 | not supplied, it defaults to `labels=['python']`.
42 |
43 | :param faktory: address of the Faktory instance to connect to.
44 | :type faktory: string
45 | :param concurrency: number of worker processes to start
46 | :type concurrency: int
47 | :param disconnect_wait: number of seconds to wait when worker is interrupted before failing all jobs
48 | :type disconnect_wait: int
49 | :param log: logger to use for status, errors and connection details
50 | :type log: logging.Logger
51 | :param labels: labels to show in the Faktory webui for this worker
52 | :type labels: tuple
53 | :param use_threads: Set to True to use threads rather than multiple processes for work to be executed on
54 | :type use_threads: bool
55 | :param executor: Set the class of the process executor that will be used. By default concurrenct.futures.ProcessPoolExecutor is used.
56 | :type executor: class
57 | """
58 | self.concurrency = kwargs.pop("concurrency", 1)
59 | self.disconnect_wait = kwargs.pop("disconnect_wait", 15)
60 | self.log = kwargs.pop("log", logging.getLogger("faktory.worker"))
61 |
62 | self._queues = kwargs.pop(
63 | "queues",
64 | [
65 | "default",
66 | ],
67 | )
68 | self._executor_class = kwargs.pop(
69 | "executor",
70 | ThreadPoolExecutor
71 | if kwargs.pop("use_threads", False)
72 | else ProcessPoolExecutor,
73 | )
74 | self._last_heartbeat = None
75 | self._tasks = dict()
76 | self._pending = list()
77 | self._disconnect_after = None
78 | self._executor = None
79 |
80 | signal.signal(signal.SIGTERM, self.handle_sigterm)
81 |
82 | if "labels" not in kwargs:
83 | kwargs["labels"] = ["python"]
84 | self.labels = kwargs["labels"]
85 |
86 | if "worker_id" not in kwargs:
87 | kwargs["worker_id"] = self.get_worker_id()
88 | self.worker_id = kwargs["worker_id"]
89 |
90 | self.faktory = Connection(*args, **kwargs)
91 | # self.faktory.debug = True
92 |
93 | def register(self, name: str, func: Callable, bind: bool = False) -> None:
94 | """
95 | Register a task that can be run with this worker.
96 |
97 | If you set bind=True, then the first argument passed to the function will always be Faktory's `jid` for this task.
98 |
99 | You can register a task after the worker has started.
100 |
101 | :param name: name of the task
102 | :type name: str
103 | :param func: function to call when the
104 | :type func: callable
105 | :param bind: pass the jid to `func`
106 | :type bind: bool
107 | :return:
108 | :rtype:
109 | """
110 | if not callable(func):
111 | raise ValueError("task func is not callable")
112 |
113 | self._tasks[name] = Task(name=name, func=func, bind=bind)
114 | self.log.info("Registered task: {}".format(name))
115 |
116 | def deregister(self, name: str) -> None:
117 | """
118 | Remove a task from the list of registered tasks.
119 |
120 | Can be called after the worker has started, any currently processing copies of `task` will continue.
121 |
122 | :param name: task name
123 | :type name: str
124 | :return:
125 | :rtype:
126 | """
127 | if name in self._tasks:
128 | del self._tasks[name]
129 | self.log.debug("Removed registered task: {}".format(name))
130 |
131 | def run(self):
132 | """
133 | Start the worker
134 |
135 | `run()` will trap signals, on the first ctrl-c it will try to gracefully shut the worker down, waiting up to 15
136 | seconds for in progress tasks to complete.
137 |
138 | If after 30 seconds tasks are still running, they are forced to terminate and the worker will close.
139 |
140 | This method is blocking -- it will only return when the worker has shutdown, either by control-c or by
141 | terminating it from the Faktory Web UI.
142 |
143 | :return:
144 | :rtype:
145 | """
146 | # create a pool of workers
147 | if not self.faktory.is_connected:
148 | self.faktory.connect(worker_id=self.worker_id)
149 |
150 | self.log.debug(
151 | "Creating a worker pool with concurrency of {}".format(self.concurrency)
152 | )
153 |
154 | self._last_heartbeat = datetime.now() + timedelta(
155 | seconds=self.send_heartbeat_every
156 | ) # schedule a heartbeat for the future
157 |
158 | self.log.info("Queues: {}".format(", ".join(self.get_queues())))
159 | self.log.info("Labels: {}".format(", ".join(self.faktory.labels)))
160 |
161 | while True:
162 | try:
163 | # tick runs continuously to process events from the faktory connection
164 | self.tick()
165 | if not self.faktory.is_connected:
166 | break
167 | except KeyboardInterrupt as e:
168 | # 1st time through: soft close, wait 15 seconds for jobs to finish and send the work results to faktory
169 | # 2nd time through: force close, don't wait, fail all current jobs and quit as quickly as possible
170 | if self.is_disconnecting:
171 | break
172 |
173 | self.log.info(
174 | "Shutdown: waiting up to 15 seconds for workers to finish current tasks"
175 | )
176 | self.disconnect(wait=self.disconnect_wait)
177 | except (BrokenProcessPool, BrokenThreadPool):
178 | self.log.info("Shutting down due to pool failure")
179 | self.disconnect(force=True, wait=15)
180 | break
181 |
182 | if self.faktory.is_connected:
183 | self.log.warning("Forcing worker processes to shutdown...")
184 | self.disconnect(force=True)
185 |
186 | self.executor.shutdown(wait=False)
187 | sys.exit(1)
188 |
189 | def disconnect(self, force=False, wait=30):
190 | """
191 | Disconnect from the Faktory server and shutdown this worker.
192 |
193 | The default is to shutdown gracefully, allowing 15s for in progress tasks to complete and update Faktory.
194 |
195 | :param force: Immediate shutdown, cancelling running tasks
196 | :type force: bool
197 | :param wait: Graceful shutdown, allowing `wait` seconds for in progress jobs to complete
198 | :type wait: int
199 | :return:
200 | :rtype:
201 | """
202 | self.log.debug(
203 | "Disconnecting from Faktory, force={} wait={}".format(force, wait)
204 | )
205 |
206 | self.is_quiet = True
207 | self.is_disconnecting = True
208 | self._disconnect_after = datetime.now() + timedelta(seconds=wait)
209 |
210 | if force:
211 | self.fail_all_jobs()
212 | self.faktory.disconnect()
213 |
214 | def tick(self):
215 | if self._pending:
216 | self.send_status_to_faktory()
217 |
218 | if self.should_send_heartbeat:
219 | self.heartbeat()
220 |
221 | if self.should_fetch_job:
222 | # grab a job to do, and start it processing
223 | job = self.faktory.fetch(self.get_queues())
224 | if job:
225 | jid = job.get("jid")
226 | func = job.get("jobtype")
227 | args = job.get("args")
228 | self._process(jid, func, args)
229 | else:
230 | if self.is_disconnecting:
231 | if self.can_disconnect:
232 | # can_disconnect returns True when there are no running tasks or pending ACK / FAILs to send
233 | # so there is no more work to send back to Faktory
234 | self.faktory.disconnect()
235 | return
236 |
237 | if datetime.now() > self._disconnect_after:
238 | self.disconnect(force=True)
239 |
240 | # faktory.fetch() blocks for 2s, but if we are not fetching jobs then we need to add a delay or this process will spin
241 | time.sleep(0.25)
242 |
243 | def send_status_to_faktory(self):
244 | for future in self._pending:
245 | if future.done():
246 | self._pending.remove(future)
247 | try:
248 | future.result(timeout=1)
249 | self._ack(future.job_id)
250 | except KeyboardInterrupt:
251 | self._fail(future.job_id)
252 | self.log.exception("Received KeyboardInterrupt! failed: {}".format(future.job_id))
253 | except Exception as e:
254 | self._fail(future.job_id, exception=e)
255 | self.log.exception("Task failed: {}".format(future.job_id))
256 |
257 | def _process(self, jid: str, job: str, args):
258 | try:
259 | task = self.get_registered_task(job)
260 | if task.bind:
261 | # pass the jid as argument 1 if the task has bind=True
262 | args = [
263 | jid,
264 | ] + args
265 |
266 | self.log.debug(
267 | "Running task: {}({})".format(
268 | task.name, ", ".join([str(x) for x in args])
269 | )
270 | )
271 | future = self.executor.submit(task.func, *args)
272 | future.job_id = jid
273 | self._pending.append(future)
274 | except BrokenExecutor as e:
275 | self._executor = None
276 | self._fail(jid, exception=e)
277 | except (KeyError, Exception) as e:
278 | self._fail(jid, exception=e)
279 |
280 | def _ack(self, jid: str):
281 | self.faktory.reply("ACK", {"jid": jid})
282 | ok = next(self.faktory.get_message())
283 |
284 | def _fail(self, jid: str, exception=None):
285 | response = {"jid": jid}
286 | if exception is not None:
287 | response["errtype"] = type(exception).__name__
288 | response["message"] = str(exception)
289 |
290 | self.faktory.reply("FAIL", response)
291 | ok = next(self.faktory.get_message())
292 |
293 | def fail_all_jobs(self):
294 | for future in self._pending:
295 | if future.done():
296 | self._ack(future.job_id)
297 | continue
298 |
299 | # force the job to fail
300 | future.cancel()
301 | self._fail(future.job_id)
302 |
303 | def handle_sigterm(self, signal, frame):
304 | raise KeyboardInterrupt
305 |
306 | @property
307 | def should_fetch_job(self) -> bool:
308 | return (
309 | not (self.is_disconnecting or self.is_quiet)
310 | and len(self._pending) < self.concurrency
311 | )
312 |
313 | @property
314 | def can_disconnect(self):
315 | return len(self._pending) == 0
316 |
317 | @property
318 | def should_send_heartbeat(self) -> bool:
319 | """
320 | Checks `self._last_heartbeat` and `self.send_heartbeat_every` to figure out of this worker needs to send a
321 | heartbeat to the Faktory server. The beat should be sent once per 60s max, and defaults to once per 15s.
322 |
323 | Chances are you don't want to override this in a subclass, but change the property `self.send_heartbeat_every`
324 | instead.
325 |
326 | :return: True if this worker should heartbeat
327 | :rtype: bool
328 | """
329 | return datetime.now() > (
330 | self._last_heartbeat + timedelta(seconds=self.send_heartbeat_every)
331 | )
332 |
333 | def heartbeat(self) -> None:
334 | """
335 | Send a heartbeat to the Faktory server so it knows this worker is still alive. This is sent every
336 | `self.send_heartbeat_every` seconds. The default is once per 15 seconds. It should not be more than 60s
337 | or Faktory will drop this worker from its active list.
338 |
339 | :return:
340 | :rtype:
341 | """
342 | self.log.debug("Sending heartbeat for worker {}".format(self.worker_id))
343 | self.faktory.reply("BEAT", {"wid": self.worker_id})
344 | ok = next(self.faktory.get_message())
345 | if "state" in ok:
346 | if "quiet" in ok:
347 | if not self.is_quiet:
348 | self.log.warning(
349 | "Faktory has quieted this worker, will not run any more tasks"
350 | )
351 | self.is_quiet = True
352 | if "terminate" in ok:
353 | if not self.is_disconnecting:
354 | self.log.warning(
355 | "Faktory has asked this worker to shutdown, will cancel any pending tasks still running 25s time"
356 | )
357 | self.disconnect(wait=25)
358 | self._last_heartbeat = datetime.now()
359 |
360 | @property
361 | def executor(self) -> Executor:
362 | """
363 | Return the concurrent.futures executor instance to use for this worker.
364 |
365 | Can be passed via the `executor` argument to `__init__` or set `use_threads=True` to use the Threaded executor.
366 |
367 | The worker will use a process based executor by default.
368 |
369 | :return: executor instance
370 | :rtype: concurrent.futures.Executor
371 | """
372 | if self._executor is None:
373 | kwargs = dict(max_workers=self.concurrency)
374 | if self._executor_class is ThreadPoolExecutor:
375 | kwargs['thread_name_prefix'] = 'Worker'
376 | self._executor = self._executor_class(**kwargs)
377 | return self._executor
378 |
379 | def get_queues(self) -> Iterable:
380 | """
381 | Returns a list of queues that this worker should be process. You can override this in a subclass to adjust the
382 | queues at runtime.
383 |
384 | :return: list of queues
385 | :rtype: list
386 | """
387 | return self._queues
388 |
389 | def get_worker_id(self) -> str:
390 | """
391 | Returns a unique ID for this worker. This method is called once, during setup of the connection. It should not
392 | change the worker_id during the lifetime of the worker.
393 |
394 | If you override this method, you should return a random string of at least 8 characters and avoid collisions
395 | with other running workers.
396 |
397 | :return: unique worker id
398 | :rtype: str
399 | """
400 | return uuid.uuid4().hex
401 |
402 | def get_registered_task(self, name: str) -> Task:
403 | try:
404 | return self._tasks[name]
405 | except KeyError:
406 | raise ValueError("'{}' is not a registered task".format(name)) from None
407 |
--------------------------------------------------------------------------------
/setup.cfg:
--------------------------------------------------------------------------------
1 | [metadata]
2 | description-file = README.md
3 |
--------------------------------------------------------------------------------
/setup.py:
--------------------------------------------------------------------------------
1 | import sys
2 |
3 | from setuptools import setup
4 |
5 | dev_requires = [
6 | "black",
7 | ]
8 |
9 | test_requires = [
10 | "pytest >= 5.0, < 6.0",
11 | ]
12 |
13 | extras = {
14 | "dev": dev_requires + test_requires,
15 | "test": test_requires,
16 | }
17 |
18 | if sys.version_info < (3, 6):
19 | extras["dev"].remove("black")
20 |
21 | extras["all_extras"] = sum(extras.values(), [])
22 |
23 | setup(
24 | name="faktory",
25 | version="1.0.0",
26 | description="Python worker for the Faktory project",
27 | extras_require=extras,
28 | python_requires=">=3.7.0",
29 | classifiers=[
30 | "Development Status :: 3 - Alpha",
31 | "License :: OSI Approved :: BSD License",
32 | "Programming Language :: Python :: 3",
33 | "Topic :: System :: Distributed Computing",
34 | ],
35 | keywords="faktory worker",
36 | url="http://github.com/cdrx/faktory_python_worker",
37 | author="Chris R",
38 | license="BSD",
39 | packages=["faktory"],
40 | zip_safe=False,
41 | )
42 |
--------------------------------------------------------------------------------
/tests/test_client.py:
--------------------------------------------------------------------------------
1 | import json
2 | from typing import Iterable
3 | from unittest.mock import MagicMock
4 |
5 | import pytest
6 | from faktory._proto import Connection
7 | from faktory.client import Client
8 |
9 |
10 | @pytest.fixture
11 | def conn() -> Connection:
12 | mock_conn = MagicMock()
13 | mock_conn.connect = MagicMock(return_value=True)
14 | mock_conn.disconnect = MagicMock(return_value=None)
15 | mock_conn.reply = MagicMock(return_value=None)
16 | # Expects an iterator
17 | mock_conn.get_message = MagicMock(return_value=iter(["OK"]))
18 |
19 | return mock_conn
20 |
21 |
22 | @pytest.fixture
23 | def client(conn) -> Client:
24 | return Client(connection=conn)
25 |
26 |
27 | def test_constructor_uses_existing_conn(conn):
28 | client = Client(connection=conn)
29 | assert client.faktory == conn
30 |
31 |
32 | def test_creates_connection_from_url():
33 | client = Client("tcp://a-server:7419")
34 | assert client.faktory.host == "a-server"
35 | assert client.faktory.port == 7419
36 |
37 |
38 | def test_connect_connects_to_server(client):
39 | assert client.is_connected == False
40 |
41 | was_successful = client.connect()
42 | assert was_successful == True
43 |
44 | client.faktory.connect.assert_called_once()
45 |
46 |
47 | def test_disconnect_disconnects_from_server(client):
48 | assert client.is_connected == False
49 |
50 | was_successful = client.connect()
51 | assert was_successful == True
52 |
53 | client.faktory.connect.assert_called_once()
54 |
55 | client.disconnect()
56 |
57 | client.faktory.disconnect.assert_called_once()
58 |
59 |
60 | def test_context_manager_no_errors(client):
61 | assert client.is_connected == False
62 |
63 | with client:
64 | client.faktory.connect.assert_called_once()
65 | assert client.is_connected == True
66 |
67 | client.faktory.disconnect.assert_called_once()
68 | assert client.is_connected == False
69 |
70 |
71 | def test_context_manager_closes_on_error(client):
72 | assert client.is_connected == False
73 |
74 | try:
75 | with client:
76 | client.faktory.connect.assert_called_once()
77 | assert client.is_connected == True
78 | raise ValueError("An error!")
79 |
80 | except ValueError:
81 | pass
82 |
83 | client.faktory.disconnect.assert_called_once()
84 | assert client.is_connected == False
85 |
86 |
87 | def test_random_job_id(client):
88 | random_id = client.random_job_id()
89 | second_random_id = client.random_job_id()
90 |
91 | assert random_id != second_random_id
92 |
93 |
94 | def test_job_id_is_serializable(client):
95 | random_id = client.random_job_id()
96 | # Will raise an error
97 | json.dumps(random_id)
98 |
99 |
100 | class TestClientQueue:
101 | def test_can_queue_job(self, client: Client):
102 | was_successful = client.queue("test", args=(1, 2))
103 | assert was_successful
104 | client.faktory.reply.assert_called_once()
105 | client.faktory.get_message.assert_called_once()
106 | assert client.faktory.reply.call_args[0][0] == "PUSH"
107 |
108 | request = client.faktory.reply.call_args[0][1]
109 | assert request["jobtype"] == "test"
110 | assert request["args"] == [1, 2]
111 | assert request["jid"] is not None
112 | assert request["queue"] == "default"
113 |
114 | def test_can_queue_job_with_generator(self, client: Client):
115 | # Split out in another test because generator gets consumed
116 | # when being cast to list, so we can't use it in parametrize
117 | was_successful = client.queue("test", args=(val for val in [1, 2]))
118 | assert was_successful
119 | client.faktory.reply.assert_called_once()
120 | client.faktory.get_message.assert_called_once()
121 | assert client.faktory.reply.call_args[0][0] == "PUSH"
122 |
123 | request = client.faktory.reply.call_args[0][1]
124 | assert request["jobtype"] == "test"
125 | assert request["args"] == [1, 2]
126 | assert request["jid"] is not None
127 | assert request["queue"] == "default"
128 |
129 | @pytest.mark.parametrize(
130 | "args",
131 | [(1, "value"), [1, "value"], {1, "value"}],
132 | )
133 | @pytest.mark.parametrize("queue", ["default", "not default"])
134 | @pytest.mark.parametrize("num_retries", [1, 2])
135 | @pytest.mark.parametrize("job_priority", [3, 4])
136 | @pytest.mark.parametrize("job_name", ["test_job", "test_job_2"])
137 | @pytest.mark.parametrize("job_id", ["test_job_id", "test_job_id_2"])
138 | def test_uses_inputs_for_job_submission(
139 | self,
140 | client: Client,
141 | job_id: str,
142 | job_name: str,
143 | job_priority: int,
144 | num_retries: int,
145 | queue: str,
146 | args: Iterable,
147 | ):
148 | was_successful = client.queue(
149 | task=job_name, args=args, queue=queue, jid=job_id, priority=job_priority
150 | )
151 | assert was_successful
152 | client.faktory.reply.assert_called_once()
153 | client.faktory.get_message.assert_called_once()
154 | assert client.faktory.reply.call_args[0][0] == "PUSH"
155 | request = client.faktory.reply.call_args[0][1]
156 | assert request["jobtype"] == job_name
157 | assert request["args"] == list(args)
158 | assert request["jid"] == job_id
159 | assert request["queue"] == queue
160 | assert request["priority"] == job_priority
161 |
162 | def test_disconnects_if_not_connected(self, client: Client):
163 | client.connect = MagicMock()
164 | client.disconnect = MagicMock()
165 | was_successful = client.queue("test", args=(1, 2))
166 |
167 | client.connect.assert_called_once()
168 | client.disconnect.assert_called_once()
169 |
170 | def test_stays_connected_if_connected(self, client: Client):
171 | client.connect()
172 | client.disconnect = MagicMock()
173 | was_successful = client.queue("test", args=(1, 2))
174 |
175 | client.disconnect.assert_not_called()
176 |
177 | def test_requires_task_name(self, client: Client):
178 | with pytest.raises(ValueError):
179 | client.queue(None, args=(1, 2))
180 |
181 | def test_requires_queue_name(self, client: Client):
182 | with pytest.raises(ValueError):
183 | client.queue("test", queue=None)
184 |
185 | def test_requires_sequence_args(self, client: Client):
186 | with pytest.raises(ValueError):
187 | client.queue("test", args="will error because not sequence")
188 |
--------------------------------------------------------------------------------
/tests/test_exceptions.py:
--------------------------------------------------------------------------------
1 | import pytest
2 |
3 | from faktory.exceptions import (
4 | FaktoryAuthenticationError,
5 | FaktoryConnectionResetError,
6 | FaktoryError,
7 | FaktoryHandshakeError,
8 | )
9 |
10 |
11 | @pytest.mark.parametrize(
12 | "exc",
13 | [FaktoryAuthenticationError, FaktoryConnectionResetError, FaktoryHandshakeError],
14 | )
15 | def test_errors_are_caught_with_connection_errors(exc: Exception):
16 | """
17 | Tests to make sure that handling a `ConnectionError` will allow
18 | us to also catch the custom, more description exceptions we
19 | have implemented.
20 | """
21 |
22 | with pytest.raises(ConnectionError):
23 | raise exc
24 |
25 |
26 | @pytest.mark.parametrize(
27 | "exc",
28 | [FaktoryAuthenticationError, FaktoryConnectionResetError, FaktoryHandshakeError],
29 | )
30 | def test_inherits_from_base_project_exception(exc: Exception):
31 | """
32 | Tests to make sure that all custom exceptions that we have
33 | defined are inheriting from our base `FaktoryError` class,
34 | so that when someone handles a `FaktoryError`, any of the
35 | subclasses will be caught appropriately.
36 | """
37 |
38 | with pytest.raises(FaktoryError):
39 | raise exc
40 |
--------------------------------------------------------------------------------
/tests/test_formatting.py:
--------------------------------------------------------------------------------
1 | import os
2 | import shlex
3 | import subprocess
4 | import sys
5 |
6 | import pytest
7 |
8 | pytestmark = pytest.mark.formatting
9 |
10 |
11 | FILE_ENDING = os.path.join("tests", "test_formatting.py")
12 | POSIX = sys.platform != "win32"
13 |
14 |
15 | @pytest.mark.skipif(sys.version_info < (3, 6), reason="Black requires Python 3.6+")
16 | def test_source_code_black_formatting():
17 | # make sure we know what working directory we're in
18 | assert __file__.endswith(FILE_ENDING)
19 |
20 | faktory_dir = os.path.dirname(os.path.dirname(__file__))
21 | result = subprocess.call(
22 | shlex.split("black --check {}".format(faktory_dir), posix=POSIX)
23 | )
24 | assert result == 0, "Faktory repo did not pass Black formatting!"
25 |
--------------------------------------------------------------------------------
/tests/test_proto.py:
--------------------------------------------------------------------------------
1 | import hashlib
2 | import io
3 | import json
4 | from typing import Any, Dict, List, Optional
5 | from unittest.mock import MagicMock, call
6 |
7 | import pytest
8 | from faktory._proto import Connection
9 | from faktory.exceptions import (
10 | FaktoryAuthenticationError,
11 | FaktoryConnectionResetError,
12 | FaktoryHandshakeError,
13 | )
14 |
15 |
16 | class TestConnectionConstructor:
17 | def test_init(self):
18 | server_url = "tcp://a-server:7419"
19 | timeout = 45
20 | buffer_size = 2048
21 | worker_id = "a custom worker id"
22 | labels = ["label 1", "label 2"]
23 |
24 | conn = Connection(
25 | faktory=server_url,
26 | timeout=timeout,
27 | buffer_size=buffer_size,
28 | worker_id=worker_id,
29 | labels=labels,
30 | )
31 | assert conn.host == "a-server"
32 | assert conn.port == 7419
33 | assert conn.password is None
34 | assert conn.timeout == timeout
35 | assert conn.buffer_size == buffer_size
36 | assert conn.worker_id == worker_id
37 | assert conn.labels == labels
38 |
39 | def test_ignores_env_if_specified(self, monkeypatch):
40 | env_value = "tcp://other-server:5000"
41 | monkeypatch.setenv("FAKTORY_URL", env_value)
42 |
43 | conn = Connection("tcp://real-server:6000")
44 |
45 | assert conn.host == "real-server"
46 | assert conn.port == 6000
47 |
48 | def test_uses_env_value_if_exists_and_not_overriden(self, monkeypatch):
49 | env_value = "tcp://other-server:5000"
50 | monkeypatch.setenv("FAKTORY_URL", env_value)
51 |
52 | conn = Connection()
53 |
54 | assert conn.host == "other-server"
55 | assert conn.port == 5000
56 |
57 | def test_has_default_if_not_provided_or_in_env(self, monkeypatch):
58 | monkeypatch.delenv("FAKTORY_URL", raising=False)
59 |
60 | conn = Connection()
61 | assert conn.host is not None
62 | assert conn.port is not None
63 |
64 | def test_uses_empty_list_when_not_given_labels(self):
65 | conn = Connection()
66 | assert conn.labels == []
67 |
68 | def test_parses_server_information_from_url(self):
69 | conn = Connection("tcp://:Password123!@localhost:7419")
70 | assert conn.host == "localhost"
71 | assert conn.password == "Password123!"
72 | assert conn.port == 7419
73 |
74 |
75 | class TestConnectionValidateConnection:
76 | def test_raises_handshake_error_on_bad_initial_connection(self, monkeypatch):
77 | payload = {"key": "value"}
78 | mocked_socket = MagicMock()
79 | mocked_get_message = MagicMock(
80 | return_value=iter(["A message that doesn't start with HI"])
81 | )
82 |
83 | conn = Connection()
84 | monkeypatch.setattr(conn, "socket", mocked_socket)
85 | monkeypatch.setattr(conn, "get_message", mocked_get_message)
86 | # Doing nothing symbolizes success
87 | with pytest.raises(FaktoryHandshakeError):
88 | conn.validate_connection()
89 |
90 | def test_raises_error_on_bad_handshake_types(self, monkeypatch):
91 | """
92 | Based on the existing source code, the requested data will
93 | have the same structure, but should raise an error on
94 | incorrectly typed values.
95 | """
96 | payload = {"v": "a non integer value"}
97 | mocked_socket = MagicMock()
98 | mocked_get_message = MagicMock(
99 | return_value=iter(["HI {}".format(json.dumps(payload))])
100 | )
101 |
102 | conn = Connection()
103 | monkeypatch.setattr(conn, "socket", mocked_socket)
104 | monkeypatch.setattr(conn, "get_message", mocked_get_message)
105 |
106 | with pytest.raises(FaktoryHandshakeError):
107 | conn.validate_connection()
108 |
109 | @pytest.mark.parametrize("server_payload", [{"v": 2}, {"v": 2, "s": str(500_000)}])
110 | def test_successful_handshake(self, monkeypatch, server_payload: Dict):
111 | """
112 | This test only covers the bare minimum to make the code
113 | succeed, and should also be fed actual values
114 | that the Faktory server produces.
115 | """
116 | host_name = "testing host"
117 | pid = 5
118 | mocked_socket = MagicMock()
119 |
120 | mocked_get_message = MagicMock(
121 | return_value=iter(["HI {}".format(json.dumps(server_payload))])
122 | )
123 | mocked_validate_handshake = MagicMock()
124 |
125 | conn = Connection(labels=["test suite"])
126 | monkeypatch.setattr(conn, "socket", mocked_socket)
127 | monkeypatch.setattr(conn, "get_message", mocked_get_message)
128 | monkeypatch.setattr(conn, "_validate_handshake", mocked_validate_handshake)
129 | # Mocked calls based on environment
130 | monkeypatch.setattr(
131 | "faktory._proto.socket.gethostname", MagicMock(return_value=host_name)
132 | )
133 | monkeypatch.setattr("faktory._proto.os.getpid", MagicMock(return_value=pid))
134 |
135 | conn.validate_connection()
136 |
137 | handshake_response = mocked_validate_handshake.call_args[0][0]
138 | assert handshake_response["hostname"] == host_name
139 | assert handshake_response["pid"] == pid
140 | assert handshake_response["labels"] == conn.labels
141 |
142 | @pytest.mark.parametrize("send_worker_id", [True, False])
143 | def test_listens_to_worker_id_flag(self, send_worker_id: bool, monkeypatch):
144 | """This test should ensure that the `send_worker_id` flag is respected."""
145 | server_payload = {"v": 2}
146 | host_name = "testing host"
147 | pid = 5
148 | worker_id = "unit test worker"
149 | mocked_socket = MagicMock()
150 |
151 | mocked_get_message = MagicMock(
152 | return_value=iter(["HI {}".format(json.dumps(server_payload))])
153 | )
154 | mocked_validate_handshake = MagicMock()
155 |
156 | conn = Connection(labels=["test suite"], worker_id=worker_id)
157 | monkeypatch.setattr(conn, "socket", mocked_socket)
158 | monkeypatch.setattr(conn, "get_message", mocked_get_message)
159 | monkeypatch.setattr(conn, "_validate_handshake", mocked_validate_handshake)
160 | # Mocked calls based on environment
161 | monkeypatch.setattr(
162 | "faktory._proto.socket.gethostname", MagicMock(return_value=host_name)
163 | )
164 | monkeypatch.setattr("faktory._proto.os.getpid", MagicMock(return_value=pid))
165 |
166 | conn.validate_connection(send_worker_id=send_worker_id)
167 |
168 | handshake_response = mocked_validate_handshake.call_args[0][0]
169 | if send_worker_id:
170 | assert "wid" in handshake_response
171 | assert handshake_response["wid"] == worker_id
172 | else:
173 | assert "wid" not in handshake_response
174 |
175 | @pytest.mark.parametrize("nonce", ["1234", None])
176 | def test_uses_nonce_if_available(self, nonce: Optional[str], monkeypatch):
177 | """
178 | This test should ensure if a nonce is present while initializing
179 | the connection to the server, its used as they key to sign the
180 | password.
181 | """
182 | server_payload = {"v": 2, "s": nonce}
183 | host_name = "testing host"
184 | pid = 5
185 | worker_id = "unit test worker"
186 | mocked_socket = MagicMock()
187 | password = "Password123!"
188 | faktory_url = "tcp://:{}@localhost:7419".format(password)
189 |
190 | mocked_get_message = MagicMock(
191 | return_value=iter(["HI {}".format(json.dumps(server_payload))])
192 | )
193 | mocked_validate_handshake = MagicMock()
194 |
195 | conn = Connection(faktory_url, labels=["test suite"])
196 | monkeypatch.setattr(conn, "socket", mocked_socket)
197 | monkeypatch.setattr(conn, "get_message", mocked_get_message)
198 | monkeypatch.setattr(conn, "_validate_handshake", mocked_validate_handshake)
199 | # Mocked calls based on environment
200 | monkeypatch.setattr(
201 | "faktory._proto.socket.gethostname", MagicMock(return_value=host_name)
202 | )
203 | monkeypatch.setattr("faktory._proto.os.getpid", MagicMock(return_value=pid))
204 |
205 | conn.validate_connection()
206 |
207 | handshake_response = mocked_validate_handshake.call_args[0][0]
208 |
209 | if not nonce:
210 | assert "pwdhash" not in handshake_response
211 | else:
212 | assert "pwdhash" in handshake_response
213 | hashed_pw = handshake_response["pwdhash"]
214 | assert hashed_pw != password
215 | expected_password = hashlib.sha256(
216 | password.encode() + nonce.encode()
217 | ).hexdigest()
218 | assert expected_password == hashed_pw
219 |
220 |
221 | class TestValidateHandshake:
222 | def test_valid_response_returns_none(self, monkeypatch):
223 | payload = {"key": "value"}
224 | mocked_socket = MagicMock()
225 | mocked_reply = MagicMock()
226 | mocked_get_message = MagicMock(return_value=iter(["OK"]))
227 |
228 | conn = Connection()
229 | monkeypatch.setattr(conn, "socket", mocked_socket)
230 | monkeypatch.setattr(conn, "reply", mocked_reply)
231 | monkeypatch.setattr(conn, "get_message", mocked_get_message)
232 | # Doing nothing symbolizes success
233 | conn._validate_handshake(payload)
234 |
235 | def test_raises_auth_error_on_bad_password(self, monkeypatch):
236 | payload = {"key": "value"}
237 | mocked_socket = MagicMock()
238 | mocked_reply = MagicMock()
239 | # Minimum error message we need to get it to raise the error
240 | err_message = "ERR: invalid password"
241 | mocked_get_message = MagicMock(return_value=iter([err_message]))
242 | mocked_close = MagicMock()
243 | mocked_socket.close = mocked_close
244 |
245 | conn = Connection()
246 | monkeypatch.setattr(conn, "socket", mocked_socket)
247 | monkeypatch.setattr(conn, "reply", mocked_reply)
248 | monkeypatch.setattr(conn, "get_message", mocked_get_message)
249 | with pytest.raises(FaktoryAuthenticationError):
250 | conn._validate_handshake(payload)
251 |
252 | mocked_close.assert_called_once()
253 |
254 | def test_raises_handshake_error_on_misc_error(self, monkeypatch):
255 | payload = {"key": "value"}
256 | mocked_socket = MagicMock()
257 | mocked_reply = MagicMock()
258 | mocked_get_message = MagicMock(
259 | return_value=iter(["ERR: A different, non auth error occurred"])
260 | )
261 |
262 | mocked_close = MagicMock()
263 | mocked_socket.close = mocked_close
264 |
265 | conn = Connection()
266 | monkeypatch.setattr(conn, "socket", mocked_socket)
267 | monkeypatch.setattr(conn, "reply", mocked_reply)
268 | monkeypatch.setattr(conn, "get_message", mocked_get_message)
269 | with pytest.raises(FaktoryHandshakeError):
270 | conn._validate_handshake(payload)
271 |
272 | mocked_close.assert_called_once()
273 |
274 | @pytest.mark.parametrize(
275 | "payload",
276 | [
277 | {"key": "value"},
278 | {"hostname": "tests", "pid": 5, "labels": ["tests"], "wid": "test worker"},
279 | {"key1": "value1", "key2": "value2"},
280 | ],
281 | )
282 | def test_uses_payload(self, monkeypatch, payload: Dict[str, Any]):
283 | mocked_socket = MagicMock()
284 | mocked_reply = MagicMock()
285 | mocked_get_message = MagicMock(return_value=iter(["OK"]))
286 |
287 | conn = Connection()
288 | monkeypatch.setattr(conn, "socket", mocked_socket)
289 | monkeypatch.setattr(conn, "reply", mocked_reply)
290 | monkeypatch.setattr(conn, "get_message", mocked_get_message)
291 |
292 | conn._validate_handshake(payload)
293 |
294 | mocked_reply.assert_called_once()
295 | mocked_get_message.assert_called_once()
296 |
297 | used_payload = mocked_reply.call_args[0][1]
298 | assert used_payload == payload
299 |
300 |
301 | class TestConnectionSelectData:
302 | def test_returns_buffer_if_exists(self, monkeypatch):
303 | mocked_socket = MagicMock()
304 |
305 | mocked_select = MagicMock(return_value=[True, False, False])
306 |
307 | socket_received_data = b"I'm the mocked out data!"
308 | mocked_recv = MagicMock(return_value=socket_received_data)
309 | mocked_socket.recv = mocked_recv
310 |
311 | conn = Connection()
312 |
313 | monkeypatch.setattr(conn, "socket", mocked_socket)
314 | monkeypatch.setattr("faktory._proto.select.select", mocked_select)
315 |
316 | data = conn.select_data(buffer_size=5_000)
317 | mocked_recv.assert_called_once()
318 | assert data == b"I'm the mocked out data!"
319 |
320 | def test_fetches_provided_buffer_size(self, monkeypatch):
321 | mocked_socket = MagicMock()
322 |
323 | mocked_select = MagicMock(return_value=[True, False, False])
324 |
325 | mocked_recv = MagicMock(return_value=b"I'm the mocked out data!")
326 | mocked_socket.recv = mocked_recv
327 |
328 | conn = Connection()
329 |
330 | monkeypatch.setattr(conn, "socket", mocked_socket)
331 | monkeypatch.setattr("faktory._proto.select.select", mocked_select)
332 |
333 | data = conn.select_data(buffer_size=5_000)
334 | mocked_recv.assert_called_once_with(5_000)
335 |
336 | def test_raises_error_on_empty_buffer(self, monkeypatch):
337 | """
338 | Tests to make sure that a "read ready" socket that
339 | has an empty buffer causes the connection to close
340 | and the connection reset to be raised.
341 | """
342 | mocked_socket = MagicMock()
343 | mocked_disconnect = MagicMock()
344 |
345 | mocked_select = MagicMock(return_value=[True, False, False])
346 |
347 | socket_received_data = b""
348 | mocked_recv = MagicMock(return_value=socket_received_data)
349 | mocked_socket.recv = mocked_recv
350 |
351 | conn = Connection()
352 |
353 | monkeypatch.setattr(conn, "socket", mocked_socket)
354 | monkeypatch.setattr("faktory._proto.select.select", mocked_select)
355 | monkeypatch.setattr(conn, "disconnect", mocked_disconnect)
356 |
357 | with pytest.raises(FaktoryConnectionResetError):
358 | data = conn.select_data(buffer_size=5_000)
359 |
360 | mocked_recv.assert_called_once()
361 | mocked_disconnect.assert_called_once()
362 |
363 | def test_raises_error_if_socket_not_ready(self, monkeypatch):
364 | mocked_socket = MagicMock()
365 | mocked_disconnect = MagicMock()
366 |
367 | mocked_select = MagicMock(return_value=[False, False, False])
368 |
369 | socket_received_data = b"I'm the mocked out data!"
370 | mocked_recv = MagicMock(return_value=socket_received_data)
371 | mocked_socket.recv = mocked_recv
372 |
373 | conn = Connection()
374 |
375 | monkeypatch.setattr(conn, "socket", mocked_socket)
376 | monkeypatch.setattr("faktory._proto.select.select", mocked_select)
377 | monkeypatch.setattr(conn, "disconnect", mocked_disconnect)
378 |
379 | with pytest.raises(FaktoryConnectionResetError):
380 | data = conn.select_data(buffer_size=5_000)
381 |
382 | mocked_recv.assert_not_called()
383 | mocked_disconnect.assert_called_once()
384 |
385 |
386 | class TestConnectionFetch:
387 | @pytest.mark.parametrize(
388 | "queues",
389 | [["default"], ["default", "important"], ["important", "not_important"]],
390 | )
391 | def test_fetches_using_queues(self, monkeypatch, queues: List[str]):
392 | mock_reply = MagicMock()
393 | mock_get_message = MagicMock(return_value=iter([None]))
394 | conn = Connection()
395 |
396 | monkeypatch.setattr(conn, "reply", mock_reply)
397 | monkeypatch.setattr(conn, "get_message", mock_get_message)
398 |
399 | conn.fetch(queues)
400 |
401 | mock_reply.assert_called_once()
402 | mock_get_message.assert_called_once()
403 |
404 | reply_args = mock_reply.call_args[0][0]
405 | assert "FETCH" in reply_args
406 | for queue in queues:
407 | assert queue in reply_args
408 |
409 | def test_returns_none_if_no_job(self, monkeypatch):
410 | mock_reply = MagicMock()
411 | mock_get_message = MagicMock(return_value=iter([None]))
412 | conn = Connection()
413 |
414 | monkeypatch.setattr(conn, "reply", mock_reply)
415 | monkeypatch.setattr(conn, "get_message", mock_get_message)
416 |
417 | queues = ["default"]
418 | result = conn.fetch(queues)
419 |
420 | mock_reply.assert_called_once()
421 | mock_get_message.assert_called_once()
422 | assert result is None
423 |
424 | def test_returns_deserialized_job_info_if_present(self, monkeypatch):
425 | data = {
426 | "wid": "test worker",
427 | "pid": 500,
428 | "labels": ["tests"],
429 | "hostname": "test suite",
430 | }
431 | serialized_data = json.dumps(data)
432 | mock_reply = MagicMock()
433 | mock_get_message = MagicMock(return_value=iter([serialized_data]))
434 | conn = Connection()
435 |
436 | monkeypatch.setattr(conn, "reply", mock_reply)
437 | monkeypatch.setattr(conn, "get_message", mock_get_message)
438 |
439 | queues = ["default"]
440 | result = conn.fetch(queues)
441 |
442 | mock_reply.assert_called_once()
443 | mock_get_message.assert_called_once()
444 | assert result == data
445 |
446 |
447 | class TestConnectionIsSupportedServerVersion:
448 | def test_supported_version(self):
449 | conn = Connection()
450 |
451 | result = conn.is_supported_server_version(2)
452 | assert result == True
453 |
454 | @pytest.mark.parametrize("version", [1, 3, 4, 5])
455 | def test_unsupported_version(self, version: int):
456 | conn = Connection()
457 |
458 | result = conn.is_supported_server_version(version)
459 | assert result == False
460 |
461 |
462 | class TestConnectionReply:
463 | @pytest.mark.parametrize("data", [None, "string data", {"data": "dict data"}])
464 | def test_sends_bytes(self, monkeypatch, data):
465 | test_command = "test command"
466 | mocked_socket = MagicMock()
467 | mock_send = MagicMock(return_value=len(test_command) + 2)
468 | mocked_socket.send = mock_send
469 |
470 | conn = Connection()
471 | monkeypatch.setattr(conn, "socket", mocked_socket)
472 |
473 | conn.reply(test_command)
474 |
475 | mock_send.assert_called_once()
476 | assert isinstance(mock_send.call_args[0][0], bytes)
477 | decoded_value = mock_send.call_args[0][0].decode()
478 | assert decoded_value == "test command\r\n"
479 |
480 | @pytest.mark.parametrize("data", [None, "string data", {"data": "dict data"}])
481 | def test_adds_return_and_newline_to_payload(self, monkeypatch, data):
482 | test_command = "test command"
483 | mocked_socket = MagicMock()
484 | mock_send = MagicMock(return_value=len(test_command) + 2)
485 | mocked_socket.send = mock_send
486 |
487 | conn = Connection()
488 | monkeypatch.setattr(conn, "socket", mocked_socket)
489 |
490 | conn.reply(test_command)
491 | called_with = mock_send.call_args[0][0]
492 | decoded = called_with.decode()
493 |
494 | assert "\r\n" == decoded[-2:]
495 |
496 | @pytest.mark.parametrize("data", [None, "string data", {"data": "dict data"}])
497 | def test_sends_until_out_of_data(self, monkeypatch, data):
498 | mocked_socket = MagicMock()
499 | mock_send = MagicMock(return_value=7)
500 | mocked_socket.send = mock_send
501 |
502 | conn = Connection()
503 | monkeypatch.setattr(conn, "socket", mocked_socket)
504 |
505 | conn.reply("test command")
506 |
507 | calls = [call(b"test command\r\n"), call(b"mmand\r\n")]
508 | mock_send.assert_has_calls(calls)
509 |
510 | @pytest.mark.parametrize("data", [None, "string data", {"data": "dict data"}])
511 | def test_send_zero_raises_error(self, monkeypatch, data):
512 | mocked_socket = MagicMock()
513 | mock_send = MagicMock(return_value=0)
514 | mocked_socket.send = mock_send
515 |
516 | conn = Connection()
517 | monkeypatch.setattr(conn, "socket", mocked_socket)
518 |
519 | with pytest.raises(FaktoryConnectionResetError):
520 | conn.reply("test command")
521 |
522 |
523 | class TestConnectionDisconnect:
524 | def test_closes_socket(self, monkeypatch):
525 | mocked_socket = MagicMock()
526 | mock_close = MagicMock()
527 | mocked_socket.close = mock_close
528 | conn = Connection()
529 | monkeypatch.setattr(conn, "socket", mocked_socket)
530 |
531 | conn.disconnect()
532 |
533 | mock_close.assert_called_once()
534 |
535 | def test_sets_connection_flag_to_false(self, monkeypatch):
536 | mocked_socket = MagicMock()
537 | conn = Connection()
538 | monkeypatch.setattr(conn, "socket", mocked_socket)
539 |
540 | conn.is_connected = True
541 | conn.disconnect()
542 |
543 | assert conn.is_connected == False
544 |
--------------------------------------------------------------------------------
/tests/test_worker.py:
--------------------------------------------------------------------------------
1 | import json
2 | import time
3 | import uuid
4 | from datetime import datetime, timedelta
5 | from typing import Any, Callable, Iterable, Optional
6 | from unittest.mock import MagicMock, PropertyMock
7 |
8 | import pytest
9 |
10 | from faktory._proto import Connection
11 | from faktory.client import Client
12 | from faktory.worker import Task, Worker
13 |
14 |
15 | @pytest.fixture
16 | def conn() -> Connection:
17 | mock_conn = MagicMock()
18 | mock_conn.connect = MagicMock(return_value=True)
19 | mock_conn.disconnect = MagicMock(return_value=None)
20 | mock_conn.reply = MagicMock(return_value=None)
21 | # Expects an iterator
22 | mock_conn.get_message = MagicMock(return_value=iter(["OK"]))
23 |
24 | return mock_conn
25 |
26 |
27 | @pytest.fixture
28 | def worker(conn: MagicMock, monkeypatch):
29 | mock_executor = MagicMock()
30 |
31 | monkeypatch.setattr("faktory.worker.Connection", conn)
32 |
33 | work = Worker()
34 | work._last_heartbeat = datetime.now()
35 | monkeypatch.setattr(work, "_executor", mock_executor)
36 | return work
37 |
38 |
39 | def create_future(
40 | done_return_value: bool = True,
41 | result_side_effect: Optional[Any] = None,
42 | result_return_value: Optional[Any] = None,
43 | ):
44 | fut = MagicMock()
45 | fut.done = MagicMock(return_value=done_return_value)
46 | fut.job_id = uuid.uuid4().hex
47 | fut.result = MagicMock(
48 | return_value=result_return_value, side_effect=result_side_effect
49 | )
50 | return fut
51 |
52 |
53 | def test_get_queues(worker):
54 | assert worker.get_queues() == worker._queues
55 |
56 |
57 | def test_mocked_connection_correctly(conn, monkeypatch):
58 | """
59 | Quick test to make sure we're mocking out the Worker's
60 | connection correctly so that we're okay to assume
61 | the rest of the test suite is valid.
62 | """
63 |
64 | conn_mock = MagicMock(return_value=conn)
65 |
66 | monkeypatch.setattr("faktory.worker.Connection", conn_mock)
67 | worker = Worker()
68 | assert worker.faktory == conn
69 |
70 |
71 | class TestWorkerTick:
72 | """
73 | Lots of tests in this class mock `time.sleep`. Some are
74 | for actual testing purposes (and will assert so), others
75 | are to make sure the test suite doesn't take forever
76 | to run.
77 | """
78 |
79 | def add_mocks(
80 | self,
81 | worker,
82 | monkeypatch,
83 | *,
84 | should_fetch_job: bool = True,
85 | should_send_heartbeat: bool = True
86 | ) -> MagicMock:
87 | """
88 | Since the mocking process on the `worker` object takes a decent
89 | number of lines of code, we're breaking it out into a helper function
90 | here. This (kind of illogically) returns the mock for `time`,
91 | since that's the only object we create here within
92 | the helper but can't access outside of it.
93 | """
94 |
95 | mock_time = MagicMock()
96 | mock_time.sleep = MagicMock()
97 |
98 | send_status_to_factory = MagicMock()
99 | heartbeat = MagicMock()
100 | process = MagicMock()
101 | disconnect = MagicMock()
102 | mock_should_fetch_job = PropertyMock(return_value=should_fetch_job)
103 | mock_should_send_heartbeat = PropertyMock(return_value=should_send_heartbeat)
104 |
105 | monkeypatch.setattr("faktory.worker.time", mock_time)
106 | monkeypatch.setattr(worker, "heartbeat", heartbeat)
107 | monkeypatch.setattr(worker, "disconnect", disconnect)
108 | monkeypatch.setattr(worker, "send_status_to_faktory", send_status_to_factory)
109 | monkeypatch.setattr(type(worker), "should_fetch_job", mock_should_fetch_job)
110 | monkeypatch.setattr(
111 | type(worker), "should_send_heartbeat", mock_should_send_heartbeat
112 | )
113 | # Mocking a "private" function might be considered a bad idea, but if we don't,
114 | # we'd need to mock out the entire success and failure paths of a succesfully
115 | # run future, which is way out of scope for unit tests.
116 | monkeypatch.setattr(worker, "_process", process)
117 |
118 | return mock_time
119 |
120 | def test_updates_faktory_on_pending(self, worker, monkeypatch):
121 | jobs = [create_future(), create_future()]
122 | mock_time = self.add_mocks(worker, monkeypatch)
123 |
124 | worker._pending = jobs
125 |
126 | worker.tick()
127 |
128 | worker.send_status_to_faktory.assert_called_once()
129 |
130 | def test_doesnt_update_when_no_pending(self, worker, monkeypatch):
131 | self.add_mocks(worker, monkeypatch)
132 |
133 | worker._pending = []
134 |
135 | worker.tick()
136 |
137 | worker.send_status_to_faktory.assert_not_called()
138 |
139 | def test_sleeps_if_not_fetching(self, worker, monkeypatch):
140 | mock_time = self.add_mocks(
141 | worker, monkeypatch, should_fetch_job=False, should_send_heartbeat=False
142 | )
143 |
144 | worker.tick()
145 |
146 | mock_time.sleep.assert_called_once()
147 |
148 | def test_submits_job_if_present(self, worker, monkeypatch):
149 | job = {
150 | "jid": "our job ID",
151 | "jobtype": "super important probably",
152 | "args": [1, 2],
153 | }
154 | self.add_mocks(worker, monkeypatch)
155 | worker.faktory.fetch.return_value = {
156 | "jid": "our job ID",
157 | "jobtype": "super important probably",
158 | "args": [1, 2],
159 | }
160 |
161 | worker.tick()
162 |
163 | worker._process.assert_called_once_with(job["jid"], job["jobtype"], job["args"])
164 |
165 | @pytest.mark.parametrize("empty_payload", [None, {}])
166 | def test_doesnt_submit_if_no_work(self, worker, monkeypatch, empty_payload):
167 | """
168 | Tests to make sure that even if we should be fetching a job,
169 | if there's no job, we don't submit anything for execution.
170 | """
171 |
172 | self.add_mocks(worker, monkeypatch)
173 | worker.faktory.fetch.return_value = empty_payload
174 |
175 | worker.tick()
176 |
177 | worker._process.assert_not_called()
178 |
179 | def test_can_disconnect_gracefully_after_emptying_jobs(self, worker, monkeypatch):
180 | worker._pending = []
181 | self.add_mocks(worker, monkeypatch, should_fetch_job=False)
182 | worker.is_disconnecting = True
183 | worker._disconnect_after = datetime.now() + timedelta(seconds=10)
184 |
185 | worker.tick()
186 |
187 | worker.faktory.disconnect.assert_called_once_with()
188 |
189 | def test_waits_on_finishing_jobs_before_forcing(self, worker, monkeypatch):
190 | """
191 | Tests the scenario between where a disconnect has started and when
192 | its forced that if there are jobs that are still pending, the
193 | worker won't kill them off until the timeout.
194 | """
195 |
196 | worker._pending = [create_future(done_return_value=False)]
197 | self.add_mocks(worker, monkeypatch, should_fetch_job=False)
198 | worker.is_disconnecting = True
199 | worker._disconnect_after = datetime.now() + timedelta(seconds=10)
200 |
201 | worker.tick()
202 | worker.faktory.disconnect.assert_not_called()
203 | worker.disconnect.assert_not_called()
204 |
205 | def test_forces_disconnect_after_timeout(self, worker, monkeypatch):
206 | worker._pending = [create_future()]
207 | self.add_mocks(worker, monkeypatch, should_fetch_job=False)
208 | worker.is_disconnecting = True
209 | worker._disconnect_after = datetime.now() - timedelta(seconds=10)
210 |
211 | worker.tick()
212 |
213 | worker.disconnect.assert_called_once_with(force=True)
214 |
215 |
216 | class TestWorkerRegister:
217 | def test_errors_without_callable_func(self, worker: Worker):
218 | with pytest.raises(ValueError):
219 | worker.register("test", "will_cause_error", True)
220 |
221 | @pytest.mark.parametrize("bind", [True, False])
222 | def test_can_register_with_and_without_binds(self, worker: Worker, bind: bool):
223 | def func(x):
224 | return x
225 |
226 | worker.register("test", func, bind=bind)
227 |
228 | def test_registers_task_instance(self, worker: Worker):
229 | def func(x):
230 | return x
231 |
232 | worker.register("test", func)
233 | task = worker.get_registered_task("test")
234 | assert isinstance(task, Task)
235 |
236 | def test_can_get_registered_task(self, worker: Worker):
237 | # This can arguably be under this set of tests
238 | # or the `get_registered_task` tests,
239 | # but it makes sure that any registered task
240 | # can be retrieved by name.
241 | def func(x):
242 | return x
243 |
244 | worker.register("test", func)
245 | task = worker.get_registered_task("test")
246 |
247 | assert task.func == func
248 | assert task.name == "test"
249 | assert task.bind == False
250 |
251 |
252 | class TestWorkerHeartbeat:
253 | def test_assigns_heartbeat_timestamp(self, worker):
254 | worker.faktory.get_message.return_value = iter(
255 | json.dumps([{"unneeded": "value"}])
256 | )
257 | worker._last_heartbeat = None
258 |
259 | now = datetime.now()
260 | time.sleep(0.001)
261 |
262 | assert worker._last_heartbeat is None
263 | worker.heartbeat()
264 | assert worker._last_heartbeat is not None
265 | assert now < worker._last_heartbeat
266 |
267 | def test_sends_heartbeat_with_worker_id(self, worker):
268 | worker.faktory.get_message.return_value = iter(
269 | json.dumps([{"unneeded": "value"}])
270 | )
271 |
272 | worker.heartbeat()
273 |
274 | heartbeat_args = worker.faktory.reply.call_args[0]
275 | command, args = heartbeat_args
276 | assert command == "BEAT"
277 | assert args == {"wid": worker.worker_id}
278 |
279 | def test_terminates_if_told(self, worker):
280 | worker.faktory.get_message.return_value = iter(
281 | [{"state": "not sure", "terminate": True}]
282 | )
283 |
284 | worker._last_heartbeat = None
285 |
286 | assert worker.is_disconnecting == False
287 | assert worker._last_heartbeat is None
288 | now = datetime.now()
289 | time.sleep(0.001)
290 | worker.heartbeat()
291 | assert worker.is_disconnecting == True
292 |
293 | assert worker._last_heartbeat is not None
294 | assert now < worker._last_heartbeat
295 |
296 | def test_quiets_if_told(self, worker):
297 | worker._last_heartbeat = None
298 | worker.is_quiet = False
299 | worker.faktory.get_message.return_value = iter(
300 | [json.dumps({"state": "not sure", "quiet": True})]
301 | )
302 |
303 | assert worker.is_quiet == False
304 | assert worker._last_heartbeat is None
305 | now = datetime.now()
306 | time.sleep(0.001)
307 | worker.heartbeat()
308 |
309 | assert worker._last_heartbeat is not None
310 | assert now < worker._last_heartbeat
311 | assert worker.is_quiet == True
312 |
313 |
314 | class TestWorkerGetWorkerId:
315 | def test_is_different(self, worker: Worker):
316 | worker_id = worker.get_worker_id()
317 | second_worker_id = worker.get_worker_id()
318 | assert worker_id != second_worker_id
319 |
320 | def test_is_json_serializable(self, worker: Worker):
321 | worker_id = worker.get_worker_id()
322 | try:
323 | json.dumps(worker_id)
324 | except TypeError as exc:
325 | pytest.fail("Worker id: {} isn't JSON serializable".format(worker_id))
326 |
327 |
328 | class TestGetRegisteredTask:
329 | def test_finds_existing_task(self, worker):
330 | def func(x):
331 | return x
332 |
333 | worker.register("test", func)
334 | task = worker.get_registered_task("test")
335 | assert task.func == func
336 | assert task.name == "test"
337 |
338 | def test_errors_on_bad_task(self, worker):
339 | def func(x):
340 | return x
341 |
342 | worker.register("test", func)
343 | with pytest.raises(ValueError):
344 | task = worker.get_registered_task("other")
345 |
346 |
347 | class TestWorkerCanDisconnect:
348 | def test_can_disconnect_if_no_tasks(self, worker):
349 | assert worker.can_disconnect is True
350 |
351 | def test_cant_disconnect_if_processing(self, worker):
352 | worker._pending.append("some item")
353 | assert worker.can_disconnect is False
354 |
355 |
356 | class TestWorkerShouldFetchJob:
357 | def test_cant_fetch_job_if_disconnecting(self, worker):
358 | worker.is_disconnecting = True
359 | assert worker.should_fetch_job is False
360 |
361 | def test_cant_fetch_job_if_quiet(self, worker):
362 | worker.is_quiet = True
363 | assert worker.should_fetch_job is False
364 |
365 | def test_cant_fetch_job_if_has_enough(self, worker):
366 | worker.is_quiet = False
367 | worker.is_disconnecting = False
368 |
369 | worker._pending.append("item")
370 | worker._pending.append("item")
371 | worker._pending.append("item")
372 |
373 | worker.concurrency = 2
374 |
375 | assert worker.should_fetch_job is False
376 |
377 | def test_can_fetch_if_needs_work(self, worker):
378 | worker.is_quiet = False
379 | worker.is_disconnecting = False
380 |
381 | worker._pending.append("item")
382 |
383 | worker.concurrency = 2
384 |
385 | assert worker.should_fetch_job is True
386 |
387 |
388 | class TestWorkerSendStatusToFaktory:
389 | def test_does_nothing_if_work_isnt_done(self, worker, monkeypatch):
390 | job_1 = create_future(done_return_value=False)
391 | job_2 = create_future(done_return_value=False)
392 |
393 | jobs = [job_1, job_2]
394 |
395 | worker._pending = jobs
396 |
397 | monkeypatch.setattr(worker, "_ack", MagicMock())
398 | monkeypatch.setattr(worker, "_fail", MagicMock())
399 |
400 | worker.send_status_to_faktory()
401 |
402 | job_1.result.assert_not_called()
403 | job_2.result.assert_not_called()
404 | worker._ack.assert_not_called()
405 | worker._fail.assert_not_called()
406 |
407 | def test_acks_successful_jobs(self, worker, monkeypatch):
408 | job_1 = create_future(done_return_value=True)
409 |
410 | worker._pending = [job_1]
411 |
412 | worker.send_status_to_faktory()
413 |
414 | job_1.result.assert_called_once()
415 | worker.faktory.reply.assert_called_with("ACK", {"jid": job_1.job_id})
416 | worker.faktory.reply.assert_called_once()
417 |
418 | worker.faktory.get_message.assert_called_once()
419 |
420 | def test_fails_errored_jobs(self, worker, monkeypatch):
421 | job_1 = create_future(
422 | done_return_value=True, result_side_effect=KeyboardInterrupt()
423 | )
424 |
425 | worker._pending = [job_1]
426 |
427 | worker.send_status_to_faktory()
428 |
429 | job_1.result.assert_called_once()
430 | worker.faktory.reply.assert_called_once_with("FAIL", {"jid": job_1.job_id})
431 |
432 | worker.faktory.get_message.assert_called_once()
433 |
434 | def test_relays_exception_info_on_errors(self, worker):
435 | job = create_future(result_side_effect=ValueError("our testing error"))
436 | worker._pending = [job]
437 | worker.send_status_to_faktory()
438 |
439 | worker.faktory.reply.assert_called_once_with(
440 | "FAIL",
441 | {
442 | "jid": job.job_id,
443 | "errtype": "ValueError",
444 | "message": "our testing error",
445 | },
446 | )
447 |
448 | worker.faktory.get_message.assert_called_once()
449 |
450 |
451 | class TestWorkerFailAllJobs:
452 | def test_doesnt_fail_successful_jobs(self, worker, monkeypatch):
453 | jobs = [create_future(), create_future(), create_future()]
454 | worker._pending = jobs
455 | monkeypatch.setattr(worker, "_ack", MagicMock())
456 | monkeypatch.setattr(worker, "_fail", MagicMock())
457 |
458 | worker.fail_all_jobs()
459 |
460 | assert worker._ack.call_count == 3
461 | worker._fail.assert_not_called()
462 |
463 | def test_fails_pending_jobs(self, worker, monkeypatch):
464 | success_job = create_future()
465 | pending_job = create_future(done_return_value=False)
466 | jobs = [success_job, pending_job]
467 |
468 | worker._pending = jobs
469 | monkeypatch.setattr(worker, "_ack", MagicMock())
470 | monkeypatch.setattr(worker, "_fail", MagicMock())
471 |
472 | worker.fail_all_jobs()
473 |
474 | worker._ack.assert_called_once_with(success_job.job_id)
475 |
476 | worker._fail.assert_called_once_with(pending_job.job_id)
477 |
478 |
479 | class TestExecutor:
480 | def test_creates_if_not_exists(self, worker):
481 | worker._executor = None
482 | worker._executor_class = MagicMock()
483 |
484 | assert worker._executor is None
485 |
486 | executor = worker.executor
487 | worker._executor_class.assert_called_once_with(max_workers=worker.concurrency)
488 | assert worker._executor is not None
489 |
490 | def test_returns_if_existing(self, worker):
491 | executor = MagicMock()
492 | worker._executor = executor
493 | worker._executor_class = MagicMock()
494 |
495 | executor = worker.executor
496 | worker._executor_class.assert_not_called()
497 |
498 | assert worker._executor == executor
499 |
500 | def test_only_creates_once(self, worker):
501 | worker._executor = None
502 | worker._executor_class = MagicMock()
503 |
504 | assert worker._executor is None
505 |
506 | executor = worker.executor
507 | other_executor = worker.executor
508 | worker._executor_class.assert_called_once_with(max_workers=worker.concurrency)
509 | assert worker._executor is not None
510 |
--------------------------------------------------------------------------------