├── .coveragerc ├── .github └── workflows │ └── ci.yml ├── .gitignore ├── .pre-commit-config.yaml ├── .travis.yml ├── CHANGELOG.rst ├── LICENSE.md ├── MANIFEST.in ├── Makefile ├── README.rst ├── development.txt ├── example ├── Dockerfile ├── README.md ├── api │ ├── example.py │ ├── helpers.py │ └── tasks.py ├── app.py ├── config │ ├── __init__.py │ ├── common.py │ └── development.py ├── env ├── requirements.txt └── supervisord.conf ├── pyproject.toml ├── pyqs ├── __init__.py ├── decorator.py ├── events.py ├── main.py ├── utils.py └── worker.py ├── setup.cfg ├── setup.py └── tests ├── __init__.py ├── tasks.py ├── test_events.py ├── test_manager_worker.py ├── test_simple_manager_worker.py ├── test_simple_worker.py ├── test_tasks.py ├── test_worker.py └── utils.py /.coveragerc: -------------------------------------------------------------------------------- 1 | [run] 2 | source=pyqs 3 | omit=*/tests/*, 4 | 5 | [report] 6 | exclude_lines = 7 | if __name__ == .__main__.: 8 | raise NotImplemented. 9 | def __repr__ 10 | 11 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | on: [push] 4 | 5 | jobs: 6 | test: 7 | runs-on: ubuntu-latest 8 | env: 9 | VIRTUAL_ENV: ignore 10 | steps: 11 | - uses: actions/checkout@v2 12 | - name: Set up Python 13 | uses: actions/setup-python@v2 14 | with: 15 | python-version: 3.9 16 | - name: Install dependencies 17 | run: | 18 | python -m pip install --upgrade pip 19 | pip install -r development.txt 20 | - name: Lint 21 | run: | 22 | pip install flake8 23 | python -m flake8 pyqs tests 24 | - name: Test 25 | run: | 26 | make test 27 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | dist/* 3 | .coverage 4 | *.pyc 5 | .build.log 6 | *.egg-info/* 7 | build/* 8 | htmlcov/* 9 | .idea/* 10 | *~ 11 | *# 12 | 13 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | - repo: git@github.com:pre-commit/pre-commit-hooks 2 | sha: cf550fcab3f12015f8676b8278b30e1a5bc10e70 3 | hooks: 4 | - id: trailing-whitespace 5 | exclude: \.html$ 6 | - id: end-of-file-fixer 7 | exclude: \.html$ 8 | - id: autopep8-wrapper 9 | args: ['-i', '--ignore=E309,E501'] 10 | - id: check-json 11 | - id: check-yaml 12 | - id: debug-statements 13 | - id: requirements-txt-fixer 14 | - id: flake8 15 | exclude: \/migrations\/ 16 | - repo: git@github.com:pre-commit/pre-commit 17 | sha: 8dba3281d5051060755459dcf88e28fc26c27526 18 | hooks: 19 | - id: validate_config 20 | - id: validate_manifest 21 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | 3 | python: 4 | - "3.6" 5 | - "3.7" 6 | - "3.8" 7 | 8 | sudo: false 9 | 10 | install: 11 | - make setup 12 | - pip install coveralls 13 | 14 | cache: 15 | directories: 16 | - $HOME/.cache/pip 17 | 18 | script: 19 | - make test 20 | 21 | after_success: 22 | - coveralls 23 | -------------------------------------------------------------------------------- /CHANGELOG.rst: -------------------------------------------------------------------------------- 1 | Changelog 2 | --------- 3 | 4 | 1.0.1 5 | ~~~~~ 6 | - Add MessageId for tracking task executions 7 | 8 | 1.0.0 9 | ~~~~~ 10 | - Drop Py2 support 11 | - Add new SimpleProcessWorker (https://github.com/spulec/PyQS/pull/76) 12 | 13 | 0.1.6 14 | ~~~~~ 15 | 16 | - Fix broken pickle of botocore clients. 17 | 18 | 0.1.5 19 | ~~~~~ 20 | 21 | - Add events hooks for pre and post processors. 22 | 23 | 0.1.4 24 | ~~~~~ 25 | 26 | - Improve behavior when a queue is created after a worker has started. The worker will now refresh the queues every 30 seconds to check for new queues. 27 | 28 | 0.1.3 29 | ~~~~~ 30 | 31 | - Change PPID checking to check for actual parent ID, instead of `PID 1`. This fixes issues running on docker containers where PPID of 1 is expected. 32 | 33 | 0.1.2 34 | ~~~~~ 35 | 36 | - 419ce2e Merge pull request #56 from orangain/honor-aws-region 37 | - 7c793d0 Merge pull request #55 from orangain/fix-indentation-error 38 | - 0643fbb Honor aws region configured by .aws/config or env var 39 | - f5c1db9 Fix indentation error 40 | - cdae257 Merge pull request #52 from cedarai/master 41 | - a2ac378 Merge pull request #53 from p1c2u/fix/nosetest-remove-stop-parameter 42 | - dbaa391 Merge pull request #51 from p1c2u/fix/pep8-styles-fixes 43 | - 1577382 Nosetest remove stop parameter 44 | - b7420e3 Add current directory to PYTHONPATH 45 | - 8d04b62 Graceful shutdown logging msg fix 46 | - 796acbc PEP8 styles fixes 47 | - 72dcb62 Merge pull request #50 from hobbsh/add_example 48 | - d00d31f Update readme 49 | - dfbf459 Use .delay() to submit messages 50 | - 612158f Merge pull request #49 from hobbsh/no_log_0_msg 51 | - 09a649f Use logger.debug for success SQS log line 52 | - dfd56c3 Fix typos in readme 53 | - a774155 Add example flask app 54 | - 17e7b7c Don't log message retrieve success when there are 0 messages 55 | - 14eb827 Add shutdown signal logging. 56 | 57 | 0.1.1 58 | ~~~~~ 59 | 60 | - Fix KeyError on accounts without queues 61 | 62 | 0.1.0 63 | ~~~~~ 64 | 65 | - Upgrade to boto3 66 | 67 | 0.0.22 68 | ~~~~~~ 69 | 70 | - Fix Python 3 support 71 | - Allow overwriting the `delay_seconds` attibute at call time 72 | 73 | 0.0.21 74 | ~~~~~~ 75 | 76 | - Add ability to tune ``PREFETCH_MULTIPLIER`` with ``--prefetch-multiplier``. 77 | 78 | 0.0.20 79 | ~~~~~~ 80 | 81 | - Respect ``--batch-size`` when sizing internal queue on ManagerWorker 82 | 83 | 0.0.19 84 | ~~~~~~ 85 | 86 | - Add ability to run with tunable BATCHSIZE and INTERVAL. Thanks to @masayang 87 | - Add ability to create tasks with a visibility delay. Thanks to @joshbuddy 88 | - Add ability to create tasks with a custom function location, allowing cross project tasks 89 | 90 | 0.0.18 91 | ~~~~~~ 92 | 93 | - Convert Changelog to .rst 94 | - Add Changelog to long description on Pypi. Thanks to @adamchainz 95 | 96 | 0.0.17 97 | ~~~~~~ 98 | 99 | - Fix typos in README 100 | - Add notes on Dead Letter Queues to README 101 | 102 | 0.0.16 103 | ~~~~~~ 104 | 105 | - Switch README to reStructuredText (.rst) format so it renders on PyPI 106 | 107 | 0.0.15 108 | ~~~~~~ 109 | 110 | - Process workers will kill themselves after attempting to process 100 111 | requests, instead of checking the internal queue 100 times. 112 | - If we find no messages on the internal queue, sleep for a moment 113 | before rechecking. 114 | 115 | 0.0.14 116 | ~~~~~~ 117 | 118 | - Process workers will kill themselves after processing 100 requests 119 | - Process workers will check a message's fetch time and visibility 120 | timeout before processing, discarding it if it has exceeded the 121 | timeout. 122 | - Log the ``process_time()`` used to process a task to the INFO level. 123 | 124 | 0.0.13 125 | ~~~~~~ 126 | 127 | - Only pass SQS Queue ID to internal queue. This is attempting to fix a 128 | bug when processing messages from multiple queues. 129 | 130 | 0.0.12 131 | ~~~~~~ 132 | 133 | - Remove extraneous debugging code 134 | 135 | 0.0.11 136 | ~~~~~~ 137 | 138 | - Add additional debugging to investigate message deletion errors 139 | 140 | 0.0.10 141 | ~~~~~~ 142 | 143 | - Give each process worker its own boto connection to avoid 144 | multiprocess race conditions during message deletion 145 | 146 | 0.0.9 147 | ----- 148 | 149 | - Change long polling interval to a valid value, 0<=LPI<=20 150 | 151 | 0.0.8 152 | ----- 153 | 154 | - Switched to long polling when pulling down messages from SQS. 155 | - Moved message deletion from SQS until after message has been 156 | processed. 157 | 158 | 0.0.7 159 | ----- 160 | 161 | - Added capability to read JSON encoded celery messages. 162 | 163 | 0.0.6 164 | ----- 165 | 166 | - Switched shutdown logging to INFO 167 | - Added brief sleep to message retrieval loop so that we don't look 168 | like we are using a ton of CPU spinning. 169 | 170 | 0.0.5 171 | ----- 172 | 173 | - Switching task failure logging to ERROR (actually this time) 174 | - Moved task success logging to INFO 175 | - Added INFO level logging for number of messages retrieved from an SQS 176 | queue. 177 | - Moved Reader and Worker process counts to DEBUG 178 | 179 | 0.0.4 180 | ----- 181 | 182 | - Added ability to pass ``region``, ``access_key_id`` and 183 | ``secret_access_key`` through to Boto when creating connections 184 | - Switched logging of task failure to the ``ERROR`` logger, from 185 | ``INFO``. 186 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2015 Steve Pulec 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include CHANGELOG.rst README.rst 2 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | PACKAGE=pyqs 2 | CUSTOM_PIP_INDEX=pypi 3 | 4 | all: setup test 5 | 6 | prepare: clean install_deps 7 | 8 | setup: prepare 9 | 10 | pre_commit: setup 11 | @pre-commit run --all-files 12 | 13 | install_deps: 14 | @if [ -z $$VIRTUAL_ENV ]; then \ 15 | echo "===================================================="; \ 16 | echo "You're not running this from a virtualenv, wtf?"; \ 17 | echo "ಠ_ಠ"; \ 18 | echo "===================================================="; \ 19 | exit 1; \ 20 | fi 21 | 22 | @if [ -z $$SKIP_DEPS ]; then \ 23 | echo "Installing missing dependencies..."; \ 24 | [ -e development.txt ] && pip install --quiet -r development.txt; \ 25 | fi 26 | @pre-commit install 27 | @python setup.py develop &> .build.log 28 | 29 | run_test: 30 | @echo "Running \033[0;32mtest suite\033[0m "; \ 31 | AWS_DEFAULT_REGION='us-east-1' nosetests --with-coverage --cover-package=$(PACKAGE) \ 32 | --cover-branches --cover-erase --verbosity=2 && pycodestyle; \ 33 | 34 | test: prepare 35 | @make run_test 36 | 37 | clean: 38 | @echo "Removing garbage..." 39 | @find . -name '*.pyc' -delete 40 | @find . -name '*.so' -delete 41 | @find . -name __pycache__ -delete 42 | @rm -rf .coverage *.egg-info *.log build dist MANIFEST yc 43 | 44 | publish: clean 45 | rm -rf dist 46 | python -m pep517.build --source --binary . 47 | twine upload dist/* 48 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | PyQS - Python task-queues for Amazon SQS |Build Status| |Coverage Status| 2 | ========================================================================= 3 | 4 | **WARNING: This library is still in beta. It has a stable API and has been deployed in production, but we have not received feedback from a large number of use cases, and it is possible there are unknown bugs.** 5 | 6 | PyQS is a simple task manager for SQS. It's goal is to provide a simple 7 | and reliable `celery `__-compatible 8 | interface to working with SQS. It uses ``boto3`` under the hood to 9 | `authenticate `__ 10 | and talk to SQS. 11 | 12 | Installation 13 | ------------ 14 | 15 | **PyQS** is available from `PyPI `__ and can 16 | be installed in all the usual ways. To install via *CLI*: 17 | 18 | .. code:: bash 19 | 20 | $ pip install pyqs 21 | 22 | Or just add it to your ``requirements.txt``. 23 | 24 | Usage 25 | ----- 26 | 27 | PyQS uses some very simple semantics to create and read tasks. Most of 28 | this comes from SQS having a very simple API. 29 | 30 | Creating Tasks 31 | ~~~~~~~~~~~~~~ 32 | 33 | Adding a task to queue is pretty simple. 34 | 35 | .. code:: python 36 | 37 | from pyqs import task 38 | 39 | @task(queue='email') 40 | def send_email(subject, message): 41 | pass 42 | 43 | send_email.delay(subject='Hi there') 44 | 45 | **NOTE:** This assumes that you have your AWS keys in the appropriate 46 | environment variables, or are using IAM roles. PyQS doesn't do anything 47 | too special to talk to AWS, it only creates the appropriate ``boto`` 48 | connection. 49 | 50 | If you don't pass a queue, PyQS will use the function path as the queue 51 | name. For example the following function lives in ``email/tasks.py``. 52 | 53 | .. code:: python 54 | 55 | @task() 56 | def send_email(subject): 57 | pass 58 | 59 | This would show up in the ``email.tasks.send_email`` queue. 60 | 61 | You can also specify the function path if you want to reference a function in a different project: 62 | 63 | .. code:: python 64 | 65 | @task(custom_function_path="foo.bar.send_email") 66 | # This references function send_email in foo/bar.py instead of email/tasks.py 67 | def send_email(subject): 68 | pass 69 | 70 | 71 | Reading Tasks 72 | ~~~~~~~~~~~~~ 73 | 74 | To read tasks we need to run PyQS. If the task is already in your 75 | ``PYTHON_PATH`` to be imported, we can just run: 76 | 77 | .. code:: bash 78 | 79 | $ pyqs email.tasks.send_email 80 | 81 | If we want to run all tasks with a certain prefix. This is based on 82 | Python's `fnmatch `__. 83 | 84 | .. code:: bash 85 | 86 | $ pyqs email.* 87 | 88 | We can also read from multiple different queues with one call by 89 | delimiting with commas: 90 | 91 | .. code:: bash 92 | 93 | $ pyqs send_email,read_email,write_email 94 | 95 | If you want to run more workers to process tasks, you can up the 96 | concurrency. This will spawn additional processes to work through 97 | messages. 98 | 99 | .. code:: bash 100 | 101 | $ pyqs send_email --concurrency 10 102 | 103 | Simple Process Worker 104 | ~~~~~~~~~~~~~~~~~~~~~ 105 | 106 | To use a simpler version of PyQS that deals with some of the edge cases in the original implementation, pass the ``simple-worker`` flag. 107 | 108 | .. code:: bash 109 | 110 | $ pyqs send_email --simple-worker 111 | 112 | The Simple Process Worker differs in the following way from the original implementation. 113 | 114 | * Does not use an internal queue and removes support for the ``prefetch-multiplier`` flag. This helps simply the mental model required, as messages are not on both the SQS queue and an internal queue. 115 | * When the ``simple-worker`` flag is passed, the default ``batchsize`` is 1 instead of 10. This is configurable. 116 | * Does not check the visibility timeout when reading or processing a message from SQS. 117 | * Allowing the worker to process the message even past its visibility timeout means we solve the problem of never processing a message if ``max_receives=1`` and we incorrectly set a shorter visibility timeout and exceed the visibility timeout. Previously, this message would have ended up in the DLQ, if one was configured, and never actually processed. 118 | * It increases the probability that we process a message more than once, especially if ``batchsize > 1``, but this can be solved by the developer checking if the message has already been processed. 119 | 120 | Hooks 121 | ~~~~~ 122 | 123 | PyQS has an event registry which can be used to run a function before or after every tasks runs. 124 | 125 | .. code:: python 126 | 127 | from pyqs import task, events 128 | 129 | def print_pre_process(context): 130 | print({"pre_process": context}) 131 | 132 | def print_post_process(context): 133 | print({"post_process": context}) 134 | 135 | events.register_event("pre_process", print_pre_process) 136 | events.register_event("post_process", print_post_process) 137 | 138 | @task(queue="my_queue") 139 | def send_email(subject): 140 | pass 141 | 142 | Operational Notes 143 | ~~~~~~~~~~~~~~~~~ 144 | 145 | **Dead Letter Queues** 146 | 147 | It is recommended to use a `Dead Letter Queue `__ 148 | for any queues that are managed by PyQS. This is because the current strategy 149 | for fetching messages does not delete them upon initial receipt. A message is 150 | **ONLY** deleted from SQS upon successful completion. **This is probably 151 | unexpected behavior if you are coming from Celery with SQS.** Celery attempted 152 | to manage this behavior internally, with varying success. 153 | 154 | If an error arises during message processing, it will be discarded and will 155 | re-appear after the visibility timeout. This can lead to behavior where 156 | there are messages that will never leave the queue and continuously throw 157 | errors. A Dead Letter Queue helps resolve this by collecting messages that 158 | have be retried a specified number of times. 159 | 160 | **Worker Seppuku** 161 | 162 | Each process worker will shut itself down after ``100`` tasks have been 163 | processed (or failed to process). This is to prevent issues with stale 164 | connections lingering and blocking tasks forever. In addition it helps 165 | guard against memory leaks, though in a rather brutish fashion. After 166 | the process worker shut itself down the managing process should notice 167 | and restart it promptly. The value of ``100`` is currently hard-coded, 168 | but could be configurable. 169 | 170 | **Queue Blocking** 171 | 172 | While there are multiple workers for reading from different queues, they 173 | all append to the same internal queue. This means that if you have one 174 | queue with lots of fast tasks, and another with a few slow tasks, they 175 | can block eachother and the fast tasks can build up behind the slow 176 | tasks. The simplest solution is to just run two different ``PyQS`` 177 | commands, one for each queue with appropriate concurrency settings. 178 | 179 | **Visibility Timeout** 180 | 181 | Care is taken to not process messages that have exceeded the visibility 182 | timeout of their queue. The goal is to prevent double processing of 183 | tasks. However, it is still quite possible for this to happen since we 184 | do not use transactional semantics around tasks. Therefore, it is 185 | important to properly set the visibility timeout on your queues based on 186 | the expected length of your tasks. If the timeout is too short, tasks 187 | will be processed twice, or very slowly. If it is too long, ephemeral 188 | failures will delay messages and reduce the queue throughput 189 | drastically. This is related to the queue blocking described above as 190 | well. SQS queues are free, so it is good practice to keep the messages 191 | stored in each as homogenous as possible. 192 | 193 | Compatibility 194 | ~~~~~~~~~~~~~ 195 | 196 | **Celery:** 197 | 198 | PyQS was created to replace celery inside of our infrastructure. To 199 | achieve this goal we wanted to make sure we were compatible with the 200 | basic Celery APIs. To this end, you can easily start trying out PyQS in 201 | your Celery-based system. PyQS can read messages that Celery has written 202 | to SQS. It will read ``pickle`` and ``json`` serialized SQS messages 203 | (Although we recommend JSON). 204 | 205 | **Operating Systems:** 206 | 207 | UNIX. Due to the use of the ``os.getppid`` system call. This feature can 208 | probably be worked around if anyone actually wants windows support. 209 | 210 | **Boto3:** 211 | 212 | Currently PyQS only supports a few basic connection parameters being 213 | explicitly passed to the connection. Any work ``boto3`` does to 214 | transparently find connection credentials, such as IAM roles, will still 215 | work properly. 216 | 217 | When running PyQS from the command-line you can pass ``--region``, 218 | ``--access-key-id``, and ``--secret-access-key`` to override the default 219 | values. 220 | 221 | Caveats 222 | ~~~~~~~ 223 | 224 | **Durability:** 225 | 226 | When we read a batch of messages from SQS we attempt to add them to our 227 | internal queue until we exceed the visibility timeout of the queue. Once 228 | this is exceeded, we discard the messages and grab a new batch. 229 | Additionally, when a process worker gets a message from the internal 230 | queue, the time the message was fetched from SQS is checked against the 231 | queues visibility timeout and discarded if it exceeds the timeout. The 232 | goal is to reduce double processing. However, this system does not 233 | provide transactions and there are cases where it is possible to process 234 | a message whos' visibility timeout has been exceeded. It is up to you to 235 | make sure that you can handle this edge case. 236 | 237 | **Task Importing:** 238 | 239 | Currently there is not advanced logic in place to find the location of 240 | modules to import tasks for processing. PyQS will try using 241 | ``importlib`` to get the module, and then find the task inside the 242 | module. Currently we wrap our usage of PyQS inside a Django admin 243 | command, which simplifies task importing. We call the 244 | `\*\*\_main()\*\* `__ 245 | method directly, skipping **main()** since it only performs argument 246 | parsing. 247 | 248 | **Running inside of containers** 249 | 250 | PyQS assumes that the process id is not 1. If you are running PyQS inside of a 251 | container, you should wrap it in supervisor or something like `dummy-init `__. 252 | 253 | **Why not just use Celery?** 254 | 255 | We like Celery. We `(Yipit.com) `__ even 256 | sponsored the `original SQS 257 | implementation `__. 258 | However, SQS is pretty different from the rest of the backends that 259 | Celery supports. Additionally the Celery team does not have the 260 | resources to create a robust SQS implementation in addition to the rest 261 | of their duties. This means the SQS is carrying around a lot extra 262 | features and a complex codebase that makes it hard to debug. 263 | 264 | We have personally experienced some very vexing resource leaks with 265 | Celery that have been hard to trackdown. For our use case, it has been 266 | simpler to switch to a simple library that we fully understand. As this 267 | library evolves that may change and the the costs of switching may not 268 | be worth it. However, we want to provide the option to others who use 269 | python and SQS to use a simpler setup. 270 | 271 | .. |Build Status| image:: https://travis-ci.org/spulec/PyQS.svg?branch=master 272 | :target: https://travis-ci.org/spulec/PyQS 273 | .. |Coverage Status| image:: https://coveralls.io/repos/spulec/PyQS/badge.svg?branch=master&service=github 274 | :target: https://coveralls.io/github/spulec/PyQS?branch=master 275 | -------------------------------------------------------------------------------- /development.txt: -------------------------------------------------------------------------------- 1 | coverage==4.4.1 2 | mock==1.0.1 3 | moto==1.3.13 4 | nose==1.3.0 5 | pep517==0.9.1 6 | pre-commit==0.7.6 7 | sure==1.2.2 8 | twine==3.3.0 9 | functools32;python_version=='2.7' 10 | pycodestyle==2.4.0 11 | -------------------------------------------------------------------------------- /example/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM python:3.6-slim-stretch 2 | 3 | # install supervisord 4 | RUN apt-get update \ 5 | && apt-get install -y supervisor libcurl4-openssl-dev gcc libssl-dev libffi-dev python3-dev curl \ 6 | && apt-get clean 7 | 8 | COPY . /api 9 | WORKDIR /api 10 | 11 | # install requirements 12 | RUN pip install -r requirements.txt 13 | 14 | # expose the app port 15 | EXPOSE 8000 16 | 17 | ENV PYTHONPATH="$PYTHONPATH:/api" 18 | 19 | # run supervisord 20 | CMD ["/usr/bin/supervisord"] 21 | -------------------------------------------------------------------------------- /example/README.md: -------------------------------------------------------------------------------- 1 | # PyQs example 2 | 3 | This example Flask app will: 4 | 5 | 1. Accept a JSON payload at `/example` 6 | 2. Publish the message (see [example.py](api/example.py)) 7 | 3. PyQS workers listening to the queue will process the messages with the specified task name 8 | 9 | Make sure you plug in your AWS API credentials to [env](env) and that a SQS queue called `queue-example` is created. 10 | 11 | ## Running with docker: 12 | 13 | * Clone this repo and cd to the root 14 | * Build image: `docker build -t pyqs-example .` 15 | * Run container: `docker run --rm --env-file env -it pyqs-example` 16 | * Post data: `docker exec -it $(docker ps | grep "pyqs-example" | awk '{print $1}') /bin/bash -c "curl -X POST -H 'Content-Type: application/json' -d '{\"message\": \"Testing\"}' http://localhost:8000/example"` 17 | 18 | ## Specifying PyQS task to be used to process your message 19 | 20 | Submitting a message to a queue is detailed in the main readme. In this example, `api/tasks.py` has a `process()` tasks that is used to process messages submitted using `tasks.process.delay(message=message_data)` in [example.py](api/example.py). See [example.py](api/example.py) and [tasks](api/tasks.py) for more information. 21 | -------------------------------------------------------------------------------- /example/api/example.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import json 3 | from flask import Blueprint, current_app, request, jsonify 4 | from api.helpers import construct_response 5 | from api import tasks 6 | 7 | blueprint = Blueprint('api', __name__) 8 | 9 | 10 | @blueprint.route('/example', methods=['POST']) 11 | def example(): 12 | queues = current_app.config.get('QUEUES') 13 | queue_name = queues['example']['name'] 14 | 15 | payload = {} 16 | data = request.get_data() 17 | 18 | try: 19 | message_data = json.loads(data) 20 | except (TypeError, AttributeError, ValueError) as e: 21 | status = 500 22 | return construct_response( 23 | 'Your payload does not appear to be valid json!', data, status) 24 | 25 | try: 26 | tasks.process.delay(message=message_data) 27 | response_message = ( 28 | f'Successfully submitted message to queue {queue_name}' 29 | ) 30 | status = 200 31 | except Exception as e: 32 | response_message = ( 33 | f'Something went wrong submitting message ' 34 | 'to queue {queue_name}! {e}' 35 | ) 36 | status = 500 37 | 38 | return construct_response(response_message, message_data, status) 39 | 40 | 41 | @blueprint.route('/health', methods=['GET']) 42 | def health(): 43 | return jsonify('OK'), 200 44 | 45 | 46 | @blueprint.route('/', methods=['GET']) 47 | def hello(): 48 | return 'Hello!' 49 | -------------------------------------------------------------------------------- /example/api/helpers.py: -------------------------------------------------------------------------------- 1 | import json 2 | import logging 3 | from flask import jsonify 4 | 5 | 6 | def construct_response(message, payload, status): 7 | body = {} 8 | 9 | if status == 500: 10 | body['message'] = ( 11 | 'Something went wrong constructing response. ' 12 | 'Is your payload valid JSON?' 13 | ) 14 | body['request_payload'] = str(payload) 15 | else: 16 | body['message'] = message 17 | body['request_payload'] = payload 18 | 19 | body['status_code'] = status 20 | logging.debug(body) 21 | 22 | resp = jsonify(body) 23 | resp.status_code = status 24 | 25 | return resp 26 | -------------------------------------------------------------------------------- /example/api/tasks.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from flask import Flask, current_app 3 | from pyqs import task 4 | 5 | 6 | # This task listens to 'queue-example' SQS queues for messages with 7 | # {"task": "process"} 8 | @task(queue='queue-example') 9 | def process(message): 10 | logging.info( 11 | f'PyQS task process() is processing message {message}. ' 12 | 'Clap your hands!' 13 | ) 14 | -------------------------------------------------------------------------------- /example/app.py: -------------------------------------------------------------------------------- 1 | import os 2 | from flask import Flask 3 | from config import config 4 | from api.example import blueprint as example_blueprint 5 | 6 | config_name = os.environ.get('ENV', 'development') 7 | current_config = config[config_name] 8 | 9 | 10 | def create_app(): 11 | app = Flask(__name__) 12 | app.config.from_object(current_config) 13 | 14 | return app 15 | 16 | 17 | app = create_app() 18 | app.register_blueprint(example_blueprint) 19 | 20 | if __name__ == "__main__": 21 | app.run(debug=True) 22 | -------------------------------------------------------------------------------- /example/config/__init__.py: -------------------------------------------------------------------------------- 1 | from .development import DevelopmentConfig 2 | 3 | config = { 4 | 'development': DevelopmentConfig 5 | # 'production': ProductionConfig 6 | } 7 | -------------------------------------------------------------------------------- /example/config/common.py: -------------------------------------------------------------------------------- 1 | import os 2 | import logging 3 | 4 | 5 | class Config(object): 6 | 7 | AWS_ACCESS_KEY_ID = os.environ.get('AWS_ACCESS_KEY_ID') 8 | AWS_SECRET_ACCESS_KEY = os.environ.get('AWS_SECRET_ACCESS_KEY') 9 | LOG_LEVEL = logging.DEBUG 10 | BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 11 | SECRET_KEY = os.environ.get('FLASK_SECRET_KEY', '') 12 | -------------------------------------------------------------------------------- /example/config/development.py: -------------------------------------------------------------------------------- 1 | from .common import Config 2 | 3 | 4 | class DevelopmentConfig(Config): 5 | DEBUG = True 6 | 7 | QUEUES = { 8 | 'default': { 9 | "name": "queue-dlq", 10 | }, 11 | 'example': { 12 | 'name': 'queue-example' 13 | } 14 | } 15 | -------------------------------------------------------------------------------- /example/env: -------------------------------------------------------------------------------- 1 | ENV=development 2 | AWS_DEFAULT_REGION=us-west-2 3 | AWS_ACCESS_KEY_ID= 4 | AWS_SECRET_ACCESS_KEY= 5 | -------------------------------------------------------------------------------- /example/requirements.txt: -------------------------------------------------------------------------------- 1 | Flask 2 | pyqs 3 | gunicorn 4 | -------------------------------------------------------------------------------- /example/supervisord.conf: -------------------------------------------------------------------------------- 1 | [supervisord] 2 | nodaemon=true 3 | 4 | [program:gunicorn] 5 | stdout_logfile=/dev/stdout 6 | stdout_logfile_maxbytes=0 7 | stderr_logfile=/dev/stderr 8 | stderr_logfile_maxbytes=0 9 | command=gunicorn -b 0.0.0.0:8000 app:app 10 | 11 | [program:pyqs] 12 | stdout_logfile=/dev/stdout 13 | stdout_logfile_maxbytes=0 14 | stderr_logfile=/dev/stderr 15 | stderr_logfile_maxbytes=0 16 | command=pyqs queue-example --concurrency 1 --batchsize 1 --log-level INFO --region %(ENV_AWS_DEFAULT_REGION)s --secret-access-key %(ENV_AWS_SECRET_ACCESS_KEY)s --access-key-id %(ENV_AWS_ACCESS_KEY_ID)s 17 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools", "wheel"] 3 | build-backend = "setuptools.build_meta" 4 | -------------------------------------------------------------------------------- /pyqs/__init__.py: -------------------------------------------------------------------------------- 1 | from .decorator import task # noqa 2 | 3 | __title__ = 'pyqs' 4 | __version__ = '1.0.1' 5 | -------------------------------------------------------------------------------- /pyqs/decorator.py: -------------------------------------------------------------------------------- 1 | import json 2 | import logging 3 | 4 | import boto3 5 | from botocore.exceptions import ClientError 6 | 7 | from .utils import get_aws_region_name, function_to_import_path 8 | 9 | logger = logging.getLogger("pyqs") 10 | 11 | 12 | def get_or_create_queue(queue_name): 13 | region_name = get_aws_region_name() 14 | sqs = boto3.resource('sqs', region_name=region_name) 15 | try: 16 | return sqs.get_queue_by_name(QueueName=queue_name) 17 | except ClientError as exc: 18 | non_existent_code = 'AWS.SimpleQueueService.NonExistentQueue' 19 | if exc.response['Error']['Code'] == non_existent_code: 20 | return sqs.create_queue(QueueName=queue_name) 21 | else: 22 | raise 23 | 24 | 25 | def task_delayer(func_to_delay, queue_name, delay_seconds=None, 26 | override=False): 27 | function_path = function_to_import_path(func_to_delay, override=override) 28 | 29 | if not queue_name: 30 | # If no queue specified, use the function_path for the queue 31 | queue_name = function_path 32 | 33 | def wrapper(*args, **kwargs): 34 | queue = get_or_create_queue(queue_name) 35 | 36 | _delay_seconds = delay_seconds 37 | if '_delay_seconds' in kwargs: 38 | _delay_seconds = kwargs['_delay_seconds'] 39 | del kwargs['_delay_seconds'] 40 | 41 | logger.info("Delaying task %s: %s, %s", function_path, args, kwargs) 42 | message_dict = { 43 | 'task': function_path, 44 | 'args': args, 45 | 'kwargs': kwargs, 46 | } 47 | 48 | message = json.dumps(message_dict) 49 | if _delay_seconds is None: 50 | _delay_seconds = 0 51 | queue.send_message(MessageBody=message, DelaySeconds=_delay_seconds) 52 | 53 | return wrapper 54 | 55 | 56 | class task(object): 57 | def __init__(self, queue=None, delay_seconds=None, 58 | custom_function_path=None): 59 | self.queue_name = queue 60 | self.delay_seconds = delay_seconds 61 | self.function_path = custom_function_path 62 | 63 | def __call__(self, *args, **kwargs): 64 | func_to_wrap = args[0] 65 | function = func_to_wrap 66 | override = False 67 | if self.function_path: 68 | override = True 69 | function = self.function_path 70 | func_to_wrap.delay = task_delayer( 71 | function, self.queue_name, self.delay_seconds, override=override) 72 | return func_to_wrap 73 | -------------------------------------------------------------------------------- /pyqs/events.py: -------------------------------------------------------------------------------- 1 | """ 2 | pyqs events registry: register callback functions on pyqs events 3 | 4 | Usage: 5 | from pyqs.events import register_event 6 | 7 | register_event("pre_process", lambda context: print(context)) 8 | """ 9 | 10 | 11 | class Events: 12 | def __init__(self): 13 | self.pre_process = [] 14 | self.post_process = [] 15 | 16 | def clear(self): 17 | self.pre_process = [] 18 | self.post_process = [] 19 | 20 | 21 | # Global singleton 22 | _EVENTS = Events() 23 | 24 | 25 | class NoEventException(Exception): 26 | pass 27 | 28 | 29 | def register_event(name, callback): 30 | if hasattr(_EVENTS, name): 31 | getattr(_EVENTS, name).append(callback) 32 | else: 33 | raise NoEventException( 34 | "{name} is not a valid pyqs event.".format(name=name) 35 | ) 36 | 37 | 38 | def get_events(): 39 | return _EVENTS 40 | 41 | 42 | def clear_events(): 43 | _EVENTS.clear() 44 | -------------------------------------------------------------------------------- /pyqs/main.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | from __future__ import unicode_literals 4 | 5 | import logging 6 | import os 7 | import sys 8 | from argparse import ArgumentParser 9 | 10 | from .worker import ManagerWorker, SimpleManagerWorker 11 | from . import __version__ 12 | 13 | logger = logging.getLogger("pyqs") 14 | 15 | SIMPLE_WORKER_DEFAULT_BATCH_SIZE = 1 16 | DEFAULT_BATCH_SIZE = 10 17 | 18 | 19 | def _set_batchsize(args): 20 | batchsize = args.batchsize 21 | if batchsize: 22 | return batchsize 23 | 24 | simple_worker = args.simple_worker 25 | if simple_worker: 26 | # Default batchsize for SimpleProcessWorker 27 | return SIMPLE_WORKER_DEFAULT_BATCH_SIZE 28 | 29 | # Default batchsize for ProcessWorker 30 | return DEFAULT_BATCH_SIZE 31 | 32 | 33 | def main(): 34 | parser = ArgumentParser(description=""" 35 | Run PyQS workers for the given queues 36 | """) 37 | parser.add_argument( 38 | "-c", 39 | "--concurrency", 40 | type=int, 41 | dest="concurrency", 42 | default=1, 43 | help='Worker concurrency', 44 | action="store", 45 | ) 46 | 47 | parser.add_argument( 48 | "queues", 49 | metavar="QUEUE_NAME", 50 | nargs="+", 51 | type=str, 52 | help='Queues to process', 53 | action="store", 54 | ) 55 | 56 | parser.add_argument( 57 | "--loglevel", 58 | "--log_level", 59 | "--log-level", 60 | dest="logging_level", 61 | type=str, 62 | default="WARN", 63 | help=( 64 | 'Set logging level. ' 65 | 'This must be one of the python default logging levels' 66 | ), 67 | action="store", 68 | ) 69 | 70 | parser.add_argument( 71 | "--access-key-id", 72 | dest="access_key_id", 73 | type=str, 74 | default=None, 75 | help='AWS_ACCESS_KEY_ID used by Boto', 76 | action="store", 77 | ) 78 | 79 | parser.add_argument( 80 | "--secret-access-key", 81 | dest="secret_access_key", 82 | type=str, 83 | default=None, 84 | help='AWS_SECRET_ACCESS_KEY used by Boto', 85 | action="store", 86 | ) 87 | 88 | parser.add_argument( 89 | "--region", 90 | dest="region", 91 | type=str, 92 | default=None, 93 | help='AWS Region to connect to SQS', 94 | action="store", 95 | ) 96 | 97 | parser.add_argument( 98 | "--endpoint-url", 99 | dest="endpoint_url", 100 | type=str, 101 | default=None, 102 | help="AWS SQS endpoint url", 103 | action="store", 104 | ) 105 | 106 | parser.add_argument( 107 | "--interval", 108 | dest="interval", 109 | type=float, 110 | default=0.0, 111 | help='Time waited by a worker after processesing a message.', 112 | action="store", 113 | ) 114 | 115 | parser.add_argument( 116 | "--batchsize", 117 | dest="batchsize", 118 | type=int, 119 | default=None, 120 | help='How many messages to download at a time from SQS.', 121 | action="store", 122 | ) 123 | 124 | parser.add_argument( 125 | "--prefetch-multiplier", 126 | dest="prefetch_multiplier", 127 | type=int, 128 | default=2, 129 | help=( 130 | 'Multiplier on the size of the internal queue ' 131 | 'for prefetching SQS messages.' 132 | ), 133 | action="store", 134 | ) 135 | 136 | parser.add_argument( 137 | '--simple-worker', 138 | dest='simple_worker', 139 | default=False, 140 | action='store_true' 141 | ) 142 | 143 | args = parser.parse_args() 144 | 145 | _main( 146 | queue_prefixes=args.queues, 147 | concurrency=args.concurrency, 148 | logging_level=args.logging_level, 149 | region=args.region, 150 | access_key_id=args.access_key_id, 151 | secret_access_key=args.secret_access_key, 152 | interval=args.interval, 153 | batchsize=_set_batchsize(args), 154 | prefetch_multiplier=args.prefetch_multiplier, 155 | simple_worker=args.simple_worker, 156 | endpoint_url=args.endpoint_url, 157 | ) 158 | 159 | 160 | def _add_cwd_to_path(): 161 | cwd = os.getcwd() 162 | if cwd not in sys.path: 163 | sys.path.insert(0, cwd) 164 | 165 | 166 | def _main(queue_prefixes, concurrency=5, logging_level="WARN", 167 | region=None, access_key_id=None, secret_access_key=None, 168 | interval=1, batchsize=DEFAULT_BATCH_SIZE, prefetch_multiplier=2, 169 | simple_worker=False, endpoint_url=None): 170 | logging.basicConfig( 171 | format="[%(levelname)s]: %(message)s", 172 | level=getattr(logging, logging_level), 173 | ) 174 | logger.info("Starting PyQS version {}".format(__version__)) 175 | 176 | if simple_worker: 177 | manager = SimpleManagerWorker( 178 | queue_prefixes, concurrency, interval, batchsize, 179 | region=region, access_key_id=access_key_id, 180 | secret_access_key=secret_access_key, 181 | endpoint_url=endpoint_url, 182 | ) 183 | else: 184 | manager = ManagerWorker( 185 | queue_prefixes, concurrency, interval, batchsize, 186 | prefetch_multiplier=prefetch_multiplier, region=region, 187 | access_key_id=access_key_id, secret_access_key=secret_access_key, 188 | endpoint_url=endpoint_url, 189 | ) 190 | 191 | _add_cwd_to_path() 192 | manager.start() 193 | manager.sleep() 194 | -------------------------------------------------------------------------------- /pyqs/utils.py: -------------------------------------------------------------------------------- 1 | import base64 2 | import json 3 | import pickle 4 | 5 | import boto3 6 | 7 | 8 | def decode_message(message): 9 | message_body = message['Body'] 10 | json_body = json.loads(message_body) 11 | if 'task' in message_body: 12 | return json_body 13 | else: 14 | # Fallback to processing celery messages 15 | return decode_celery_message(json_body['body']) 16 | 17 | 18 | def decode_celery_message(json_task): 19 | message = base64.b64decode(json_task) 20 | try: 21 | return json.loads(message) 22 | except ValueError: 23 | pass 24 | return pickle.loads(message) 25 | 26 | 27 | def function_to_import_path(function, override=False): 28 | if override: 29 | return function 30 | return "{}.{}".format(function.__module__, function.__name__) 31 | 32 | 33 | def get_aws_region_name(): 34 | region_name = boto3.session.Session().region_name 35 | if not region_name: 36 | region_name = 'us-east-1' 37 | 38 | return region_name 39 | -------------------------------------------------------------------------------- /pyqs/worker.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | from __future__ import unicode_literals 3 | 4 | import copy 5 | import fnmatch 6 | import importlib 7 | import logging 8 | import os 9 | import signal 10 | import sys 11 | import traceback 12 | import time 13 | 14 | from multiprocessing import Event, Process, Queue 15 | try: 16 | from queue import Empty, Full 17 | except ImportError: 18 | from Queue import Empty, Full 19 | 20 | import boto3 21 | 22 | from pyqs.utils import get_aws_region_name, decode_message 23 | from pyqs.events import get_events 24 | 25 | MESSAGE_DOWNLOAD_BATCH_SIZE = 10 26 | LONG_POLLING_INTERVAL = 20 27 | logger = logging.getLogger("pyqs") 28 | 29 | 30 | def get_conn( 31 | region=None, access_key_id=None, secret_access_key=None, endpoint_url=None 32 | ): 33 | kwargs = { 34 | "aws_access_key_id": access_key_id, 35 | "aws_secret_access_key": secret_access_key, 36 | "region_name": region, 37 | } 38 | 39 | if endpoint_url: 40 | kwargs["endpoint_url"] = endpoint_url 41 | if not kwargs["region_name"]: 42 | kwargs["region_name"] = get_aws_region_name() 43 | 44 | return boto3.client( 45 | "sqs", 46 | **kwargs, 47 | ) 48 | 49 | 50 | class BaseWorker(Process): 51 | def __init__(self, *args, **kwargs): 52 | self._connection = None 53 | self.parent_id = kwargs.pop('parent_id') 54 | super(BaseWorker, self).__init__(*args, **kwargs) 55 | self.should_exit = Event() 56 | 57 | def _get_connection(self): 58 | if self._connection: 59 | return self._connection 60 | 61 | if self.connection_args is None: 62 | self._connection = get_conn() 63 | else: 64 | self._connection = get_conn(**self.connection_args) 65 | return self._connection 66 | 67 | def shutdown(self): 68 | logger.info( 69 | "Received shutdown signal, shutting down PID {}!".format( 70 | os.getpid())) 71 | self.should_exit.set() 72 | 73 | def parent_is_alive(self): 74 | if os.getppid() != self.parent_id: 75 | logger.info( 76 | "Parent process has gone away, exiting process {}!".format( 77 | os.getpid())) 78 | return False 79 | return True 80 | 81 | 82 | class ReadWorker(BaseWorker): 83 | 84 | def __init__(self, queue_url, internal_queue, batchsize, 85 | connection_args=None, *args, **kwargs): 86 | super(ReadWorker, self).__init__(*args, **kwargs) 87 | if connection_args is None: 88 | connection_args = {} 89 | self.connection_args = connection_args 90 | self.queue_url = queue_url 91 | 92 | sqs_queue = get_conn(**self.connection_args).get_queue_attributes( 93 | QueueUrl=queue_url, AttributeNames=['All'])['Attributes'] 94 | self.visibility_timeout = int(sqs_queue['VisibilityTimeout']) 95 | 96 | self.internal_queue = internal_queue 97 | self.batchsize = batchsize 98 | 99 | def run(self): 100 | # Set the child process to not receive any keyboard interrupts 101 | signal.signal(signal.SIGINT, signal.SIG_IGN) 102 | 103 | logger.info( 104 | "Running ReadWorker: {}, pid: {}".format( 105 | self.queue_url, os.getpid())) 106 | while not self.should_exit.is_set() and self.parent_is_alive(): 107 | self.read_message() 108 | self.internal_queue.close() 109 | self.internal_queue.cancel_join_thread() 110 | 111 | def read_message(self): 112 | messages = self._get_connection().receive_message( 113 | QueueUrl=self.queue_url, 114 | MaxNumberOfMessages=self.batchsize, 115 | WaitTimeSeconds=LONG_POLLING_INTERVAL, 116 | ).get('Messages', []) 117 | 118 | logger.debug( 119 | "Successfully got {} messages from SQS queue {}".format( 120 | len(messages), self.queue_url)) # noqa 121 | 122 | start = time.time() 123 | for message in messages: 124 | end = time.time() 125 | if int(end - start) >= self.visibility_timeout: 126 | # Don't add any more messages since they have 127 | # re-appeared in the sqs queue Instead just reset and get 128 | # fresh messages from the sqs queue 129 | msg = ( 130 | "Clearing Local messages since we exceeded " 131 | "their visibility_timeout" 132 | ) 133 | logger.warning(msg) 134 | break 135 | 136 | message_body = decode_message(message) 137 | try: 138 | packed_message = { 139 | "queue": self.queue_url, 140 | "message": message, 141 | "start_time": start, 142 | "timeout": self.visibility_timeout, 143 | } 144 | self.internal_queue.put( 145 | packed_message, True, self.visibility_timeout) 146 | except Full: 147 | msg = ( 148 | "Timed out trying to add the following message " 149 | "to the internal queue after {} seconds: {}" 150 | ).format(self.visibility_timeout, message_body) # noqa 151 | logger.warning(msg) 152 | continue 153 | else: 154 | logger.debug( 155 | "Message successfully added to internal queue " 156 | "from SQS queue {} with body: {}".format( 157 | self.queue_url, message_body)) # noqa 158 | 159 | 160 | class BaseProcessWorker(BaseWorker): 161 | def __init__(self, *args, **kwargs): 162 | super(BaseProcessWorker, self).__init__(*args, **kwargs) 163 | 164 | def _run_hooks(self, hook_name, context): 165 | hooks = getattr(get_events(), hook_name) 166 | for hook in hooks: 167 | hook(context) 168 | 169 | def _create_pre_process_context(self, packed_message): 170 | message = packed_message['message'] 171 | message_body = decode_message(message) 172 | full_task_path = message_body['task'] 173 | 174 | pre_process_context = { 175 | "message_id": message['MessageId'], 176 | "task_name": full_task_path.split(".")[-1], 177 | "args": message_body['args'], 178 | "kwargs": message_body['kwargs'], 179 | "full_task_path": full_task_path, 180 | "fetch_time": packed_message['start_time'], 181 | "queue_url": packed_message['queue'], 182 | "timeout": packed_message['timeout'], 183 | "receipt_handle": message['ReceiptHandle'] 184 | } 185 | 186 | return pre_process_context 187 | 188 | def _get_task(self, full_task_path): 189 | 190 | task_name = full_task_path.split(".")[-1] 191 | task_path = ".".join(full_task_path.split(".")[:-1]) 192 | task_module = importlib.import_module(task_path) 193 | task = getattr(task_module, task_name) 194 | 195 | return task 196 | 197 | def _process_task(self, pre_process_context): 198 | task = self._get_task(pre_process_context["full_task_path"]) 199 | 200 | # Modify the contexts separately so the original 201 | # context isn't modified by later processing 202 | post_process_context = copy.copy(pre_process_context) 203 | 204 | start_time = time.time() 205 | try: 206 | self._run_hooks("pre_process", pre_process_context) 207 | task(*pre_process_context["args"], **pre_process_context["kwargs"]) 208 | except Exception: 209 | end_time = time.time() 210 | logger.exception( 211 | "Task {} raised error in {:.4f} seconds: with args: {} " 212 | "and kwargs: {}: {}".format( 213 | pre_process_context["full_task_path"], 214 | end_time - start_time, 215 | pre_process_context["args"], 216 | pre_process_context["kwargs"], 217 | traceback.format_exc(), 218 | ) 219 | ) 220 | post_process_context["status"] = "exception" 221 | post_process_context["exception"] = traceback.format_exc() 222 | self._run_hooks("post_process", post_process_context) 223 | return True 224 | else: 225 | end_time = time.time() 226 | self._get_connection().delete_message( 227 | QueueUrl=pre_process_context["queue_url"], 228 | ReceiptHandle=pre_process_context["receipt_handle"] 229 | ) 230 | logger.info( 231 | "Processed task {} in {:.4f} seconds with args: {} " 232 | "and kwargs: {}".format( 233 | pre_process_context["full_task_path"], 234 | end_time - start_time, 235 | pre_process_context["args"], 236 | pre_process_context["kwargs"], 237 | ) 238 | ) 239 | post_process_context["status"] = "success" 240 | self._run_hooks("post_process", post_process_context) 241 | return True 242 | 243 | 244 | class ProcessWorker(BaseProcessWorker): 245 | 246 | def __init__(self, internal_queue, interval, connection_args=None, *args, 247 | **kwargs): 248 | super(ProcessWorker, self).__init__(*args, **kwargs) 249 | self.connection_args = connection_args 250 | self.internal_queue = internal_queue 251 | self.interval = interval 252 | self._messages_to_process_before_shutdown = 100 253 | self.messages_processed = 0 254 | 255 | def run(self): 256 | # Set the child process to not receive any keyboard interrupts 257 | signal.signal(signal.SIGINT, signal.SIG_IGN) 258 | 259 | logger.info("Running ProcessWorker, pid: {}".format(os.getpid())) 260 | 261 | while not self.should_exit.is_set() and self.parent_is_alive(): 262 | processed = self.process_message() 263 | if processed: 264 | self.messages_processed += 1 265 | time.sleep(self.interval) 266 | else: 267 | # If we have no messages wait a moment before rechecking. 268 | time.sleep(0.001) 269 | if self.messages_processed \ 270 | >= self._messages_to_process_before_shutdown: 271 | self.shutdown() 272 | 273 | def process_message(self): 274 | try: 275 | packed_message = self.internal_queue.get(timeout=0.5) 276 | except Empty: 277 | # Return False if we did not attempt to process any messages 278 | return False 279 | 280 | pre_process_context = self._create_pre_process_context(packed_message) 281 | 282 | current_time = time.time() 283 | if int(current_time - pre_process_context["fetch_time"]) \ 284 | >= pre_process_context["timeout"]: 285 | logger.warning( 286 | "Discarding task {} with args: {} and kwargs: {} due to " 287 | "exceeding visibility timeout".format( # noqa 288 | pre_process_context["full_task_path"], 289 | repr(pre_process_context["args"]), 290 | repr(pre_process_context["kwargs"]), 291 | ) 292 | ) 293 | return True 294 | 295 | return self._process_task(pre_process_context) 296 | 297 | 298 | class SimpleProcessWorker(BaseProcessWorker): 299 | 300 | def __init__(self, queue_url, interval, batchsize, 301 | connection_args=None, *args, **kwargs): 302 | super(SimpleProcessWorker, self).__init__(*args, **kwargs) 303 | if connection_args is None: 304 | connection_args = {} 305 | self.connection_args = connection_args 306 | self.queue_url = queue_url 307 | 308 | sqs_queue = get_conn(**self.connection_args).get_queue_attributes( 309 | QueueUrl=queue_url, AttributeNames=['All'])['Attributes'] 310 | self.visibility_timeout = int(sqs_queue['VisibilityTimeout']) 311 | 312 | self.interval = interval 313 | self.batchsize = batchsize 314 | self._messages_to_process_before_shutdown = 100 315 | self.messages_processed = 0 316 | 317 | def run(self): 318 | # Set the child process to not receive any keyboard interrupts 319 | signal.signal(signal.SIGINT, signal.SIG_IGN) 320 | 321 | logger.info( 322 | "Running SimpleProcessWorker: {}, pid: {}".format( 323 | self.queue_url, os.getpid())) 324 | 325 | while not self.should_exit.is_set() and self.parent_is_alive(): 326 | messages = self.read_message() 327 | start = time.time() 328 | 329 | for message in messages: 330 | packed_message = { 331 | "queue": self.queue_url, 332 | "message": message, 333 | "start_time": start, 334 | "timeout": self.visibility_timeout, 335 | } 336 | 337 | processed = self.process_message(packed_message) 338 | 339 | if processed: 340 | self.messages_processed += 1 341 | time.sleep(self.interval) 342 | 343 | if self.messages_processed \ 344 | >= self._messages_to_process_before_shutdown: 345 | self.shutdown() 346 | 347 | def read_message(self): 348 | messages = self._get_connection().receive_message( 349 | QueueUrl=self.queue_url, 350 | MaxNumberOfMessages=self.batchsize, 351 | WaitTimeSeconds=LONG_POLLING_INTERVAL, 352 | ).get('Messages', []) 353 | 354 | logger.debug( 355 | "Successfully got {} messages from SQS queue {}".format( 356 | len(messages), self.queue_url)) # noqa 357 | 358 | return messages 359 | 360 | def process_message(self, packed_message): 361 | 362 | pre_process_context = self._create_pre_process_context(packed_message) 363 | 364 | return self._process_task(pre_process_context) 365 | 366 | 367 | class BaseManager(object): 368 | 369 | def __init__(self, queue_prefixes, interval, batchsize, 370 | region=None, access_key_id=None, 371 | secret_access_key=None, endpoint_url=None): 372 | self.connection_args = { 373 | "region": region, 374 | "access_key_id": access_key_id, 375 | "secret_access_key": secret_access_key, 376 | "endpoint_url": endpoint_url, 377 | } 378 | self.interval = interval 379 | self.batchsize = batchsize 380 | if batchsize > MESSAGE_DOWNLOAD_BATCH_SIZE: 381 | self.batchsize = MESSAGE_DOWNLOAD_BATCH_SIZE 382 | if batchsize <= 0: 383 | self.batchsize = 1 384 | self.queue_prefixes = queue_prefixes 385 | self.queue_urls = self.get_queue_urls_from_queue_prefixes( 386 | self.queue_prefixes) 387 | self._pid = os.getpid() 388 | self._running = True 389 | self._register_signals() 390 | 391 | def _register_signals(self): 392 | for SIG in [signal.SIGINT, signal.SIGTERM, signal.SIGQUIT, 393 | signal.SIGHUP]: 394 | self.register_shutdown_signal(SIG) 395 | 396 | def get_queue_urls_from_queue_prefixes(self, queue_prefixes): 397 | conn = get_conn(**self.connection_args) 398 | queue_urls = conn.list_queues().get('QueueUrls', []) 399 | matching_urls = [] 400 | 401 | logger.info("Loading Queues:") 402 | for prefix in queue_prefixes: 403 | logger.info("[Queue]\t{}".format(prefix)) 404 | matching_urls.extend([ 405 | queue_url for queue_url in queue_urls if 406 | fnmatch.fnmatch(queue_url.rsplit("/", 1)[1], prefix) 407 | ]) 408 | logger.info("Found matching SQS Queues: {}".format(matching_urls)) 409 | return matching_urls 410 | 411 | def check_for_new_queues(self): 412 | raise NotImplementedError 413 | 414 | def start(self): 415 | raise NotImplementedError 416 | 417 | def stop(self): 418 | raise NotImplementedError 419 | 420 | def sleep(self): 421 | counter = 0 422 | while self._running: 423 | counter = counter + 1 424 | if counter % 1000 == 0: 425 | self.process_counts() 426 | self.replace_workers() 427 | if counter % 30000 == 0: 428 | counter = 0 429 | self.check_for_new_queues() 430 | time.sleep(0.001) 431 | self._exit() 432 | 433 | def register_shutdown_signal(self, SIG): 434 | signal.signal(SIG, self._graceful_shutdown) 435 | 436 | def _graceful_shutdown(self, signum, frame): 437 | logger.info('Received shutdown signal %s', signum) 438 | self._running = False 439 | 440 | def _exit(self): 441 | logger.info('Graceful shutdown. Sending shutdown signal to children.') 442 | self.stop() 443 | sys.exit(0) 444 | 445 | def process_counts(self): 446 | raise NotImplementedError 447 | 448 | def replace_workers(self): 449 | raise NotImplementedError 450 | 451 | 452 | class SimpleManagerWorker(BaseManager): 453 | WORKER_CHILDREN_CLASS = SimpleProcessWorker 454 | 455 | def __init__(self, queue_prefixes, worker_concurrency, interval, batchsize, 456 | region=None, access_key_id=None, secret_access_key=None, 457 | endpoint_url=None): 458 | 459 | super(SimpleManagerWorker, self).__init__(queue_prefixes, interval, 460 | batchsize, region, 461 | access_key_id, 462 | secret_access_key, 463 | endpoint_url) 464 | 465 | self.worker_children = [] 466 | self._initialize_worker_children(worker_concurrency) 467 | 468 | def _initialize_worker_children(self, number): 469 | for queue_url in self.queue_urls: 470 | for index in range(number): 471 | self.worker_children.append( 472 | self.WORKER_CHILDREN_CLASS( 473 | queue_url, self.interval, self.batchsize, 474 | connection_args=self.connection_args, 475 | parent_id=self._pid, 476 | ) 477 | ) 478 | 479 | def check_for_new_queues(self): 480 | queue_urls = self.get_queue_urls_from_queue_prefixes( 481 | self.queue_prefixes) 482 | new_queue_urls = set(queue_urls) - set(self.queue_urls) 483 | for new_queue_url in new_queue_urls: 484 | logger.info("Found new queue\t{}".format(new_queue_url)) 485 | worker = self.WORKER_CHILDREN_CLASS( 486 | new_queue_url, self.interval, self.batchsize, 487 | connection_args=self.connection_args, 488 | parent_id=self._pid, 489 | ) 490 | worker.start() 491 | self.worker_children.append(worker) 492 | 493 | def start(self): 494 | for child in self.worker_children: 495 | child.start() 496 | 497 | def stop(self): 498 | for child in self.worker_children: 499 | child.shutdown() 500 | for child in self.worker_children: 501 | child.join() 502 | 503 | def process_counts(self): 504 | worker_count = sum(map(lambda x: x.is_alive(), self.worker_children)) 505 | logger.debug("Worker Processes: {}".format(worker_count)) 506 | 507 | def replace_workers(self): 508 | for index, worker in enumerate(self.worker_children): 509 | if not worker.is_alive(): 510 | logger.info( 511 | "Worker Process {} is no longer responding, " 512 | "spawning a new worker.".format(worker.pid)) 513 | self.worker_children.pop(index) 514 | worker = self.WORKER_CHILDREN_CLASS( 515 | worker.queue_url, self.interval, self.batchsize, 516 | connection_args=self.connection_args, 517 | parent_id=self._pid, 518 | ) 519 | worker.start() 520 | self.worker_children.append(worker) 521 | 522 | 523 | class ManagerWorker(BaseManager): 524 | WORKER_CHILDREN_CLASS = ProcessWorker 525 | 526 | def __init__(self, queue_prefixes, worker_concurrency, interval, batchsize, 527 | prefetch_multiplier=2, region=None, access_key_id=None, 528 | secret_access_key=None, endpoint_url=None): 529 | 530 | super(ManagerWorker, self).__init__(queue_prefixes, interval, 531 | batchsize, region, 532 | access_key_id, 533 | secret_access_key, 534 | endpoint_url) 535 | 536 | self.prefetch_multiplier = prefetch_multiplier 537 | self.worker_children = [] 538 | self.reader_children = [] 539 | 540 | self.setup_internal_queue(worker_concurrency) 541 | self._initialize_reader_children() 542 | self._initialize_worker_children(worker_concurrency) 543 | 544 | def _initialize_reader_children(self): 545 | for queue_url in self.queue_urls: 546 | self.reader_children.append( 547 | ReadWorker( 548 | queue_url, self.internal_queue, self.batchsize, 549 | connection_args=self.connection_args, 550 | parent_id=self._pid, 551 | ) 552 | ) 553 | 554 | def _initialize_worker_children(self, number): 555 | for index in range(number): 556 | self.worker_children.append( 557 | self.WORKER_CHILDREN_CLASS( 558 | self.internal_queue, self.interval, 559 | connection_args=self.connection_args, 560 | parent_id=self._pid, 561 | ) 562 | ) 563 | 564 | def check_for_new_queues(self): 565 | queue_urls = self.get_queue_urls_from_queue_prefixes( 566 | self.queue_prefixes) 567 | new_queue_urls = set(queue_urls) - set(self.queue_urls) 568 | for new_queue_url in new_queue_urls: 569 | logger.info("Found new queue\t{}".format(new_queue_url)) 570 | worker = ReadWorker( 571 | new_queue_url, self.internal_queue, self.batchsize, 572 | connection_args=self.connection_args, 573 | parent_id=self._pid, 574 | ) 575 | worker.start() 576 | self.reader_children.append(worker) 577 | 578 | def setup_internal_queue(self, worker_concurrency): 579 | self.internal_queue = Queue( 580 | worker_concurrency * self.prefetch_multiplier * self.batchsize) 581 | 582 | def start(self): 583 | for child in self.reader_children: 584 | child.start() 585 | for child in self.worker_children: 586 | child.start() 587 | 588 | def stop(self): 589 | for child in self.reader_children: 590 | child.shutdown() 591 | for child in self.reader_children: 592 | child.join() 593 | 594 | for child in self.worker_children: 595 | child.shutdown() 596 | for child in self.worker_children: 597 | child.join() 598 | 599 | def process_counts(self): 600 | reader_count = sum(map(lambda x: x.is_alive(), self.reader_children)) 601 | worker_count = sum(map(lambda x: x.is_alive(), self.worker_children)) 602 | logger.debug("Reader Processes: {}".format(reader_count)) 603 | logger.debug("Worker Processes: {}".format(worker_count)) 604 | 605 | def replace_workers(self): 606 | self._replace_reader_children() 607 | self._replace_worker_children() 608 | 609 | def _replace_reader_children(self): 610 | for index, reader in enumerate(self.reader_children): 611 | if not reader.is_alive(): 612 | logger.info( 613 | "Reader Process {} is no longer responding, " 614 | "spawning a new reader.".format(reader.pid)) 615 | queue_url = reader.queue_url 616 | self.reader_children.pop(index) 617 | worker = ReadWorker( 618 | queue_url, self.internal_queue, self.batchsize, 619 | connection_args=self.connection_args, 620 | parent_id=self._pid, 621 | ) 622 | worker.start() 623 | self.reader_children.append(worker) 624 | 625 | def _replace_worker_children(self): 626 | for index, worker in enumerate(self.worker_children): 627 | if not worker.is_alive(): 628 | logger.info( 629 | "Worker Process {} is no longer responding, " 630 | "spawning a new worker.".format(worker.pid)) 631 | self.worker_children.pop(index) 632 | worker = self.WORKER_CHILDREN_CLASS( 633 | self.internal_queue, self.interval, 634 | connection_args=self.connection_args, 635 | parent_id=self._pid, 636 | ) 637 | worker.start() 638 | self.worker_children.append(worker) 639 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [flake8] 2 | ignore = E501,E502,W293,E121,E123,E124,E125,E126,E127,E128,E265,E266 -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | from __future__ import unicode_literals 4 | 5 | import re 6 | 7 | from setuptools import setup, find_packages 8 | 9 | with open('pyqs/__init__.py', 'r') as fd: 10 | version = re.search(r'^__version__\s*=\s*[\'"]([^\'"]*)[\'"]', 11 | fd.read(), re.MULTILINE).group(1) 12 | 13 | if not version: 14 | raise RuntimeError('Cannot find version information') 15 | 16 | 17 | with open('README.rst') as readme_file: 18 | readme = readme_file.read() 19 | 20 | with open('CHANGELOG.rst') as changelog_file: 21 | changelog = changelog_file.read() 22 | 23 | setup( 24 | name='pyqs', 25 | version=version, 26 | description='A simple task-queue for SQS.', 27 | long_description=readme + '\n\n' + changelog, 28 | author='Steve Pulec', 29 | author_email='spulec@gmail.com', 30 | url='https://github.com/spulec/pyqs', 31 | entry_points={ 32 | 'console_scripts': [ 33 | 'pyqs = pyqs.main:main', 34 | ], 35 | }, 36 | install_requires=[ 37 | 'boto3>=1.7.0' 38 | ], 39 | packages=[n for n in find_packages() if not n.startswith('tests')], 40 | include_package_data=True, 41 | ) 42 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | import sure # noqa: F401 2 | -------------------------------------------------------------------------------- /tests/tasks.py: -------------------------------------------------------------------------------- 1 | from pyqs import task 2 | 3 | try: 4 | basestring 5 | except NameError: 6 | basestring = str 7 | 8 | task_results = [] 9 | 10 | 11 | @task() 12 | def index_incrementer(message, extra=None): 13 | if isinstance(message, basestring): 14 | task_results.append(message) 15 | else: 16 | raise ValueError( 17 | "Need to be given basestring, was given {}".format(message)) 18 | 19 | 20 | @task() 21 | def sleeper(message, extra=None): 22 | import time 23 | time.sleep(message) 24 | 25 | 26 | @task(queue='email') 27 | def send_email(subject, message): 28 | pass 29 | 30 | 31 | @task(queue='delayed', delay_seconds=5) 32 | def delayed_task(): 33 | pass 34 | 35 | 36 | @task(custom_function_path="custom_function.path", queue="foobar") 37 | def custom_path_task(): 38 | pass 39 | 40 | 41 | @task() 42 | def exception_task(message, extra=None): 43 | raise Exception('this task raises an exception!') 44 | -------------------------------------------------------------------------------- /tests/test_events.py: -------------------------------------------------------------------------------- 1 | from nose.tools import assert_raises 2 | from pyqs import events 3 | from tests.utils import clear_events_registry 4 | 5 | 6 | @clear_events_registry 7 | def test_register_event(): 8 | def print_pre_process(context): 9 | print(context) 10 | 11 | events.register_event("pre_process", print_pre_process) 12 | events.get_events().pre_process.should.equal([print_pre_process]) 13 | 14 | 15 | @clear_events_registry 16 | def test_register_multiple_same_events(): 17 | def print_pre_process(context): 18 | print(context) 19 | 20 | def print_numbers(context): 21 | print(1 + 2) 22 | 23 | events.register_event("pre_process", print_pre_process) 24 | events.register_event("pre_process", print_numbers) 25 | events.get_events().pre_process.should.equal([ 26 | print_pre_process, print_numbers 27 | ]) 28 | 29 | 30 | @clear_events_registry 31 | def test_register_different_events(): 32 | def print_pre_process(context): 33 | print(context) 34 | 35 | def print_post_process(context): 36 | print(context) 37 | 38 | events.register_event("pre_process", print_pre_process) 39 | events.register_event("post_process", print_post_process) 40 | events.get_events().pre_process.should.equal([print_pre_process]) 41 | events.get_events().post_process.should.equal([print_post_process]) 42 | 43 | 44 | @clear_events_registry 45 | def test_register_multiple_different_events(): 46 | def print_pre_process(context): 47 | print(context) 48 | 49 | def print_post_process(context): 50 | print(context) 51 | 52 | def print_numbers(context): 53 | print(1 + 2) 54 | 55 | events.register_event("pre_process", print_pre_process) 56 | events.register_event("pre_process", print_numbers) 57 | events.register_event("post_process", print_post_process) 58 | events.register_event("post_process", print_numbers) 59 | events.get_events().pre_process.should.equal([ 60 | print_pre_process, print_numbers 61 | ]) 62 | events.get_events().post_process.should.equal([ 63 | print_post_process, print_numbers 64 | ]) 65 | 66 | 67 | @clear_events_registry 68 | def test_register_non_existent_event(): 69 | non_existent_event = "non_existent_event" 70 | assert_raises( 71 | events.NoEventException, 72 | events.register_event, 73 | non_existent_event, 74 | lambda x: x 75 | ) 76 | -------------------------------------------------------------------------------- /tests/test_manager_worker.py: -------------------------------------------------------------------------------- 1 | import json 2 | import logging 3 | import os 4 | import signal 5 | import time 6 | 7 | import boto3 8 | from mock import patch, Mock, MagicMock 9 | from moto import mock_sqs 10 | 11 | from pyqs.main import main, _main 12 | from pyqs.worker import ManagerWorker 13 | from tests.utils import ( 14 | MockLoggingHandler, ThreadWithReturnValue2, ThreadWithReturnValue3, 15 | ) 16 | 17 | 18 | @mock_sqs 19 | def test_manager_worker_create_proper_children_workers(): 20 | """ 21 | Test managing process creates multiple child workers 22 | """ 23 | conn = boto3.client('sqs', region_name='us-east-1') 24 | conn.create_queue(QueueName="email") 25 | 26 | manager = ManagerWorker( 27 | queue_prefixes=['email'], worker_concurrency=3, interval=2, 28 | batchsize=10, 29 | ) 30 | 31 | len(manager.reader_children).should.equal(1) 32 | len(manager.worker_children).should.equal(3) 33 | 34 | 35 | @mock_sqs 36 | def test_manager_worker_with_queue_prefix(): 37 | """ 38 | Test managing process can find queues by prefix 39 | """ 40 | conn = boto3.client('sqs', region_name='us-east-1') 41 | conn.create_queue(QueueName="email.foobar") 42 | conn.create_queue(QueueName="email.baz") 43 | 44 | manager = ManagerWorker( 45 | queue_prefixes=['email.*'], worker_concurrency=1, interval=1, 46 | batchsize=10, 47 | ) 48 | 49 | len(manager.reader_children).should.equal(2) 50 | children = manager.reader_children 51 | # Pull all the read children and sort by name to make testing easier 52 | sorted_children = sorted(children, key=lambda child: child.queue_url) 53 | 54 | sorted_children[0].queue_url.should.equal( 55 | "https://queue.amazonaws.com/123456789012/email.baz") 56 | sorted_children[1].queue_url.should.equal( 57 | "https://queue.amazonaws.com/123456789012/email.foobar") 58 | 59 | 60 | @mock_sqs 61 | def test_manager_start_and_stop(): 62 | """ 63 | Test managing process can start and stop child processes 64 | """ 65 | conn = boto3.client('sqs', region_name='us-east-1') 66 | conn.create_queue(QueueName="email") 67 | 68 | manager = ManagerWorker( 69 | queue_prefixes=['email'], worker_concurrency=2, interval=1, 70 | batchsize=10, 71 | ) 72 | 73 | len(manager.worker_children).should.equal(2) 74 | 75 | manager.worker_children[0].is_alive().should.equal(False) 76 | manager.worker_children[1].is_alive().should.equal(False) 77 | 78 | manager.start() 79 | 80 | manager.worker_children[0].is_alive().should.equal(True) 81 | manager.worker_children[1].is_alive().should.equal(True) 82 | 83 | manager.stop() 84 | 85 | manager.worker_children[0].is_alive().should.equal(False) 86 | manager.worker_children[1].is_alive().should.equal(False) 87 | 88 | 89 | @patch("pyqs.main.ManagerWorker") 90 | @mock_sqs 91 | def test_main_method(ManagerWorker): 92 | """ 93 | Test creation of manager process from _main method 94 | """ 95 | _main(["email1", "email2"], concurrency=2) 96 | 97 | ManagerWorker.assert_called_once_with( 98 | ['email1', 'email2'], 2, 1, 10, prefetch_multiplier=2, 99 | region=None, secret_access_key=None, access_key_id=None, 100 | endpoint_url=None, 101 | ) 102 | ManagerWorker.return_value.start.assert_called_once_with() 103 | 104 | 105 | @patch("pyqs.main._main") 106 | @patch("pyqs.main.ArgumentParser") 107 | @mock_sqs 108 | def test_real_main_method(ArgumentParser, _main): 109 | """ 110 | Test parsing of arguments from main method 111 | """ 112 | ArgumentParser.return_value.parse_args.return_value = Mock( 113 | concurrency=3, queues=["email1"], interval=1, batchsize=5, 114 | logging_level="WARN", region='us-east-1', prefetch_multiplier=2, 115 | access_key_id=None, secret_access_key=None, simple_worker=False, 116 | endpoint_url=None, 117 | ) 118 | main() 119 | 120 | _main.assert_called_once_with( 121 | queue_prefixes=['email1'], concurrency=3, interval=1, batchsize=5, 122 | logging_level="WARN", region='us-east-1', prefetch_multiplier=2, 123 | access_key_id=None, secret_access_key=None, simple_worker=False, 124 | endpoint_url=None, 125 | ) 126 | 127 | 128 | @patch("pyqs.main._main") 129 | @patch("pyqs.main.ArgumentParser") 130 | @mock_sqs 131 | def test_real_main_method_default_batchsize(ArgumentParser, _main): 132 | """ 133 | Test parsing of arguments from main method batch default 134 | """ 135 | ArgumentParser.return_value.parse_args.return_value = Mock( 136 | concurrency=3, queues=["email1"], interval=1, batchsize=None, 137 | logging_level="WARN", region='us-east-1', prefetch_multiplier=2, 138 | access_key_id=None, secret_access_key=None, simple_worker=False, 139 | endpoint_url=None, 140 | ) 141 | main() 142 | 143 | _main.assert_called_once_with( 144 | queue_prefixes=['email1'], concurrency=3, interval=1, batchsize=10, 145 | logging_level="WARN", region='us-east-1', prefetch_multiplier=2, 146 | access_key_id=None, secret_access_key=None, simple_worker=False, 147 | endpoint_url=None, 148 | ) 149 | 150 | 151 | @mock_sqs 152 | def test_master_spawns_worker_processes(): 153 | """ 154 | Test managing process creates child workers 155 | """ 156 | 157 | # Setup SQS Queue 158 | conn = boto3.client('sqs', region_name='us-east-1') 159 | conn.create_queue(QueueName="tester") 160 | 161 | # Setup Manager 162 | manager = ManagerWorker(["tester"], 1, 1, 10) 163 | manager.start() 164 | 165 | # Check Workers 166 | len(manager.reader_children).should.equal(1) 167 | len(manager.worker_children).should.equal(1) 168 | 169 | manager.reader_children[0].is_alive().should.be.true 170 | manager.worker_children[0].is_alive().should.be.true 171 | 172 | # Cleanup 173 | manager.stop() 174 | 175 | 176 | @patch("pyqs.worker.LONG_POLLING_INTERVAL", 1) 177 | @mock_sqs 178 | def test_master_replaces_reader_processes(): 179 | """ 180 | Test managing process replaces reader children 181 | """ 182 | 183 | # Setup SQS Queue 184 | conn = boto3.client('sqs', region_name='us-east-1') 185 | conn.create_queue(QueueName="tester") 186 | 187 | # Setup Manager 188 | manager = ManagerWorker( 189 | queue_prefixes=["tester"], worker_concurrency=1, interval=1, 190 | batchsize=10, 191 | ) 192 | manager.start() 193 | 194 | # Get Reader PID 195 | pid = manager.reader_children[0].pid 196 | 197 | # Kill Reader and wait to replace 198 | manager.reader_children[0].shutdown() 199 | time.sleep(2) 200 | manager.replace_workers() 201 | 202 | # Check Replacement 203 | manager.reader_children[0].pid.shouldnt.equal(pid) 204 | 205 | # Cleanup 206 | manager.stop() 207 | 208 | 209 | @mock_sqs 210 | def test_master_counts_processes(): 211 | """ 212 | Test managing process counts child processes 213 | """ 214 | 215 | # Setup Logging 216 | logger = logging.getLogger("pyqs") 217 | del logger.handlers[:] 218 | logger.handlers.append(MockLoggingHandler()) 219 | 220 | # Setup SQS Queue 221 | conn = boto3.client('sqs', region_name='us-east-1') 222 | conn.create_queue(QueueName="tester") 223 | 224 | # Setup Manager 225 | manager = ManagerWorker(["tester"], 2, 1, 10) 226 | manager.start() 227 | 228 | # Check Workers 229 | manager.process_counts() 230 | 231 | # Cleanup 232 | manager.stop() 233 | 234 | # Check messages 235 | msg1 = "Reader Processes: 1" 236 | logger.handlers[0].messages['debug'][-2].lower().should.contain( 237 | msg1.lower()) 238 | msg2 = "Worker Processes: 2" 239 | logger.handlers[0].messages['debug'][-1].lower().should.contain( 240 | msg2.lower()) 241 | 242 | 243 | @mock_sqs 244 | def test_master_replaces_worker_processes(): 245 | """ 246 | Test managing process replaces worker processes 247 | """ 248 | # Setup SQS Queue 249 | conn = boto3.client('sqs', region_name='us-east-1') 250 | conn.create_queue(QueueName="tester") 251 | 252 | # Setup Manager 253 | manager = ManagerWorker( 254 | queue_prefixes=["tester"], worker_concurrency=1, interval=1, 255 | batchsize=10, 256 | ) 257 | manager.start() 258 | 259 | # Get Worker PID 260 | pid = manager.worker_children[0].pid 261 | 262 | # Kill Worker and wait to replace 263 | manager.worker_children[0].shutdown() 264 | time.sleep(0.1) 265 | manager.replace_workers() 266 | 267 | # Check Replacement 268 | manager.worker_children[0].pid.shouldnt.equal(pid) 269 | 270 | # Cleanup 271 | manager.stop() 272 | 273 | 274 | @mock_sqs 275 | @patch("pyqs.worker.sys") 276 | def test_master_handles_signals(sys): 277 | """ 278 | Test managing process handles OS signals 279 | """ 280 | 281 | # Setup SQS Queue 282 | conn = boto3.client('sqs', region_name='us-east-1') 283 | conn.create_queue(QueueName="tester") 284 | 285 | # Mock out sys.exit 286 | sys.exit = Mock() 287 | 288 | # Have our inner method send our signal 289 | def process_counts(): 290 | os.kill(os.getpid(), signal.SIGTERM) 291 | 292 | # Setup Manager 293 | manager = ManagerWorker( 294 | queue_prefixes=["tester"], worker_concurrency=1, interval=1, 295 | batchsize=10, 296 | ) 297 | manager.process_counts = process_counts 298 | manager._graceful_shutdown = MagicMock() 299 | 300 | # When we start and trigger a signal 301 | manager.start() 302 | manager.sleep() 303 | 304 | # Then we exit 305 | sys.exit.assert_called_once_with(0) 306 | 307 | 308 | @patch("pyqs.worker.LONG_POLLING_INTERVAL", 3) 309 | @mock_sqs 310 | def test_master_shuts_down_busy_read_workers(): 311 | """ 312 | Test managing process properly cleans up busy Reader Workers 313 | """ 314 | # For debugging test 315 | import sys 316 | logger = logging.getLogger("pyqs") 317 | logger.setLevel(logging.DEBUG) 318 | stdout_handler = logging.StreamHandler(sys.stdout) 319 | logger.addHandler(stdout_handler) 320 | 321 | # Setup SQS Queue 322 | conn = boto3.client('sqs', region_name='us-east-1') 323 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 324 | 325 | # Add Slow tasks 326 | message = json.dumps({ 327 | 'task': 'tests.tasks.sleeper', 328 | 'args': [], 329 | 'kwargs': { 330 | 'message': 5, 331 | }, 332 | }) 333 | 334 | # Fill the queue (we need a lot of messages to trigger the bug) 335 | for _ in range(20): 336 | conn.send_message(QueueUrl=queue_url, MessageBody=message) 337 | 338 | # Create function to watch and kill stuck processes 339 | def sleep_and_kill(pid): 340 | import os 341 | import signal 342 | import time 343 | # This sleep time is long enoug for 100 messages in queue 344 | time.sleep(5) 345 | try: 346 | os.kill(pid, signal.SIGKILL) 347 | except OSError: 348 | # Return that we didn't need to kill the process 349 | return True 350 | else: 351 | # Return that we needed to kill the process 352 | return False 353 | 354 | # Setup Manager 355 | manager = ManagerWorker( 356 | queue_prefixes=["tester"], worker_concurrency=0, interval=0.0, 357 | batchsize=1, 358 | ) 359 | manager.start() 360 | 361 | # Give our processes a moment to start 362 | time.sleep(1) 363 | 364 | # Setup Threading watcher 365 | try: 366 | # Try Python 2 Style 367 | thread = ThreadWithReturnValue2( 368 | target=sleep_and_kill, args=(manager.reader_children[0].pid,)) 369 | thread.daemon = True 370 | except TypeError: 371 | # Use Python 3 Style 372 | thread = ThreadWithReturnValue3( 373 | target=sleep_and_kill, args=(manager.reader_children[0].pid,), 374 | daemon=True, 375 | ) 376 | 377 | thread.start() 378 | 379 | # Stop the Master Process 380 | manager.stop() 381 | 382 | # Check if we had to kill the Reader Worker or it exited gracefully 383 | return_value = thread.join() 384 | if not return_value: 385 | raise Exception("Reader Worker failed to quit!") 386 | 387 | 388 | @mock_sqs 389 | def test_master_shuts_down_busy_process_workers(): 390 | """ 391 | Test managing process properly cleans up busy Process Workers 392 | """ 393 | # For debugging test 394 | import sys 395 | logger = logging.getLogger("pyqs") 396 | logger.setLevel(logging.DEBUG) 397 | stdout_handler = logging.StreamHandler(sys.stdout) 398 | logger.addHandler(stdout_handler) 399 | 400 | # Setup SQS Queue 401 | conn = boto3.client('sqs', region_name='us-east-1') 402 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 403 | 404 | # Add Slow tasks 405 | message = json.dumps({ 406 | 'task': 'tests.tasks.sleeper', 407 | 'args': [], 408 | 'kwargs': { 409 | 'message': 5, 410 | }, 411 | }) 412 | 413 | # Fill the queue (we need a lot of messages to trigger the bug) 414 | for _ in range(20): 415 | conn.send_message(QueueUrl=queue_url, MessageBody=message) 416 | 417 | # Create function to watch and kill stuck processes 418 | def sleep_and_kill(pid): 419 | import os 420 | import signal 421 | import time 422 | # This sleep time is long enoug for 100 messages in queue 423 | time.sleep(5) 424 | try: 425 | os.kill(pid, signal.SIGKILL) 426 | except OSError: 427 | # Return that we didn't need to kill the process 428 | return True 429 | else: 430 | # Return that we needed to kill the process 431 | return False 432 | 433 | # Setup Manager 434 | manager = ManagerWorker( 435 | queue_prefixes=["tester"], worker_concurrency=1, interval=0.0, 436 | batchsize=1, 437 | ) 438 | manager.start() 439 | 440 | # Give our processes a moment to start 441 | time.sleep(1) 442 | 443 | # Setup Threading watcher 444 | try: 445 | # Try Python 2 Style 446 | thread = ThreadWithReturnValue2( 447 | target=sleep_and_kill, args=(manager.reader_children[0].pid,)) 448 | thread.daemon = True 449 | except TypeError: 450 | # Use Python 3 Style 451 | thread = ThreadWithReturnValue3( 452 | target=sleep_and_kill, args=(manager.reader_children[0].pid,), 453 | daemon=True, 454 | ) 455 | 456 | thread.start() 457 | 458 | # Stop the Master Process 459 | manager.stop() 460 | 461 | # Check if we had to kill the Reader Worker or it exited gracefully 462 | return_value = thread.join() 463 | if not return_value: 464 | raise Exception("Reader Worker failed to quit!") 465 | 466 | 467 | @mock_sqs 468 | def test_manager_picks_up_new_queues(): 469 | """ 470 | Test that the manager will recognize new SQS queues have been added 471 | """ 472 | 473 | # Setup SQS Queue 474 | conn = boto3.client('sqs', region_name='us-east-1') 475 | 476 | # Setup Manager 477 | manager = ManagerWorker( 478 | queue_prefixes=["tester"], worker_concurrency=1, interval=1, 479 | batchsize=10, 480 | ) 481 | manager.start() 482 | 483 | # No queues found 484 | len(manager.reader_children).should.equal(0) 485 | 486 | # Create the queue 487 | conn.create_queue(QueueName="tester") 488 | manager.check_for_new_queues() 489 | 490 | # The manager should have seen the new queue was created and add a reader 491 | len(manager.reader_children).should.equal(1) 492 | manager.reader_children[0].queue_url.should.equal( 493 | "https://queue.amazonaws.com/123456789012/tester") 494 | 495 | # Cleanup 496 | manager.stop() 497 | -------------------------------------------------------------------------------- /tests/test_simple_manager_worker.py: -------------------------------------------------------------------------------- 1 | import json 2 | import logging 3 | import os 4 | import signal 5 | import time 6 | 7 | import boto3 8 | from mock import patch, Mock, MagicMock 9 | from moto import mock_sqs 10 | 11 | from pyqs.main import main, _main 12 | from pyqs.worker import SimpleManagerWorker 13 | from tests.utils import ( 14 | MockLoggingHandler, ThreadWithReturnValue2, ThreadWithReturnValue3, 15 | ) 16 | 17 | 18 | @mock_sqs 19 | def test_simple_manager_worker_create_proper_children_workers(): 20 | """ 21 | Test simple managing process creates multiple child workers 22 | """ 23 | conn = boto3.client('sqs', region_name='us-east-1') 24 | conn.create_queue(QueueName="email") 25 | 26 | manager = SimpleManagerWorker( 27 | queue_prefixes=['email'], worker_concurrency=3, interval=2, 28 | batchsize=10, 29 | ) 30 | 31 | len(manager.worker_children).should.equal(3) 32 | 33 | 34 | @mock_sqs 35 | def test_simple_manager_worker_with_queue_prefix(): 36 | """ 37 | Test simple managing process can find queues by prefix 38 | """ 39 | conn = boto3.client('sqs', region_name='us-east-1') 40 | conn.create_queue(QueueName="email.foobar") 41 | conn.create_queue(QueueName="email.baz") 42 | 43 | manager = SimpleManagerWorker( 44 | queue_prefixes=['email.*'], worker_concurrency=1, interval=1, 45 | batchsize=10, 46 | ) 47 | 48 | len(manager.worker_children).should.equal(2) 49 | children = manager.worker_children 50 | # Pull all the read children and sort by name to make testing easier 51 | sorted_children = sorted(children, key=lambda child: child.queue_url) 52 | 53 | sorted_children[0].queue_url.should.equal( 54 | "https://queue.amazonaws.com/123456789012/email.baz") 55 | sorted_children[1].queue_url.should.equal( 56 | "https://queue.amazonaws.com/123456789012/email.foobar") 57 | 58 | 59 | @mock_sqs 60 | def test_simple_manager_start_and_stop(): 61 | """ 62 | Test simple managing process can start and stop child processes 63 | """ 64 | conn = boto3.client('sqs', region_name='us-east-1') 65 | conn.create_queue(QueueName="email") 66 | 67 | manager = SimpleManagerWorker( 68 | queue_prefixes=['email'], worker_concurrency=2, interval=1, 69 | batchsize=10, 70 | ) 71 | 72 | len(manager.worker_children).should.equal(2) 73 | 74 | manager.worker_children[0].is_alive().should.equal(False) 75 | manager.worker_children[1].is_alive().should.equal(False) 76 | 77 | manager.start() 78 | 79 | manager.worker_children[0].is_alive().should.equal(True) 80 | manager.worker_children[1].is_alive().should.equal(True) 81 | 82 | manager.stop() 83 | 84 | manager.worker_children[0].is_alive().should.equal(False) 85 | manager.worker_children[1].is_alive().should.equal(False) 86 | 87 | 88 | @patch("pyqs.main.SimpleManagerWorker") 89 | @mock_sqs 90 | def test_main_method(SimpleManagerWorker): 91 | """ 92 | Test creation of simple manager process from _main method 93 | """ 94 | _main(["email1", "email2"], concurrency=2, simple_worker=True) 95 | 96 | SimpleManagerWorker.assert_called_once_with( 97 | ['email1', 'email2'], 2, 1, 10, 98 | region=None, secret_access_key=None, access_key_id=None, 99 | endpoint_url=None, 100 | ) 101 | SimpleManagerWorker.return_value.start.assert_called_once_with() 102 | 103 | 104 | @patch("pyqs.main._main") 105 | @patch("pyqs.main.ArgumentParser") 106 | @mock_sqs 107 | def test_real_main_method(ArgumentParser, _main): 108 | """ 109 | Test parsing of arguments from main method 110 | """ 111 | ArgumentParser.return_value.parse_args.return_value = Mock( 112 | concurrency=3, queues=["email1"], interval=1, batchsize=5, 113 | logging_level="WARN", region='us-east-1', prefetch_multiplier=2, 114 | access_key_id=None, secret_access_key=None, simple_worker=True, 115 | endpoint_url=None, 116 | ) 117 | main() 118 | 119 | _main.assert_called_once_with( 120 | queue_prefixes=['email1'], concurrency=3, interval=1, batchsize=5, 121 | logging_level="WARN", region='us-east-1', prefetch_multiplier=2, 122 | access_key_id=None, secret_access_key=None, simple_worker=True, 123 | endpoint_url=None, 124 | ) 125 | 126 | 127 | @patch("pyqs.main._main") 128 | @patch("pyqs.main.ArgumentParser") 129 | @mock_sqs 130 | def test_real_main_method_default_batchsize(ArgumentParser, _main): 131 | """ 132 | Test parsing of arguments from main method batch default 133 | """ 134 | ArgumentParser.return_value.parse_args.return_value = Mock( 135 | concurrency=3, queues=["email1"], interval=1, batchsize=None, 136 | logging_level="WARN", region='us-east-1', prefetch_multiplier=2, 137 | access_key_id=None, secret_access_key=None, simple_worker=True, 138 | endpoint_url=None, 139 | ) 140 | main() 141 | 142 | _main.assert_called_once_with( 143 | queue_prefixes=['email1'], concurrency=3, interval=1, batchsize=1, 144 | logging_level="WARN", region='us-east-1', prefetch_multiplier=2, 145 | access_key_id=None, secret_access_key=None, simple_worker=True, 146 | endpoint_url=None, 147 | ) 148 | 149 | 150 | @mock_sqs 151 | def test_master_spawns_worker_processes(): 152 | """ 153 | Test simple managing process creates child workers 154 | """ 155 | 156 | # Setup SQS Queue 157 | conn = boto3.client('sqs', region_name='us-east-1') 158 | conn.create_queue(QueueName="tester") 159 | 160 | # Setup Manager 161 | manager = SimpleManagerWorker(["tester"], 1, 1, 10) 162 | manager.start() 163 | 164 | # Check Workers 165 | len(manager.worker_children).should.equal(1) 166 | 167 | manager.worker_children[0].is_alive().should.be.true 168 | 169 | # Cleanup 170 | manager.stop() 171 | 172 | 173 | @mock_sqs 174 | def test_master_counts_processes(): 175 | """ 176 | Test simple managing process counts child processes 177 | """ 178 | 179 | # Setup Logging 180 | logger = logging.getLogger("pyqs") 181 | del logger.handlers[:] 182 | logger.handlers.append(MockLoggingHandler()) 183 | 184 | # Setup SQS Queue 185 | conn = boto3.client('sqs', region_name='us-east-1') 186 | conn.create_queue(QueueName="tester") 187 | 188 | # Setup Manager 189 | manager = SimpleManagerWorker(["tester"], 2, 1, 10) 190 | manager.start() 191 | 192 | # Check Workers 193 | manager.process_counts() 194 | 195 | # Cleanup 196 | manager.stop() 197 | 198 | # Check messages 199 | msg2 = "Worker Processes: 2" 200 | logger.handlers[0].messages['debug'][-1].lower().should.contain( 201 | msg2.lower()) 202 | 203 | 204 | @mock_sqs 205 | def test_master_replaces_worker_processes(): 206 | """ 207 | Test simple managing process replaces worker processes 208 | """ 209 | # Setup SQS Queue 210 | conn = boto3.client('sqs', region_name='us-east-1') 211 | conn.create_queue(QueueName="tester") 212 | 213 | # Setup Manager 214 | manager = SimpleManagerWorker( 215 | queue_prefixes=["tester"], worker_concurrency=1, interval=1, 216 | batchsize=10, 217 | ) 218 | manager.start() 219 | 220 | # Get Worker PID 221 | pid = manager.worker_children[0].pid 222 | 223 | # Kill Worker and wait to replace 224 | manager.worker_children[0].shutdown() 225 | time.sleep(0.1) 226 | manager.replace_workers() 227 | 228 | # Check Replacement 229 | manager.worker_children[0].pid.shouldnt.equal(pid) 230 | 231 | # Cleanup 232 | manager.stop() 233 | 234 | 235 | @mock_sqs 236 | @patch("pyqs.worker.sys") 237 | def test_master_handles_signals(sys): 238 | """ 239 | Test simple managing process handles OS signals 240 | """ 241 | 242 | # Setup SQS Queue 243 | conn = boto3.client('sqs', region_name='us-east-1') 244 | conn.create_queue(QueueName="tester") 245 | 246 | # Mock out sys.exit 247 | sys.exit = Mock() 248 | 249 | # Have our inner method send our signal 250 | def process_counts(): 251 | os.kill(os.getpid(), signal.SIGTERM) 252 | 253 | # Setup Manager 254 | manager = SimpleManagerWorker( 255 | queue_prefixes=["tester"], worker_concurrency=1, interval=1, 256 | batchsize=10, 257 | ) 258 | manager.process_counts = process_counts 259 | manager._graceful_shutdown = MagicMock() 260 | 261 | # When we start and trigger a signal 262 | manager.start() 263 | manager.sleep() 264 | 265 | # Then we exit 266 | sys.exit.assert_called_once_with(0) 267 | 268 | 269 | @mock_sqs 270 | def test_master_shuts_down_busy_process_workers(): 271 | """ 272 | Test simple managing process properly cleans up busy Process Workers 273 | """ 274 | # For debugging test 275 | import sys 276 | logger = logging.getLogger("pyqs") 277 | logger.setLevel(logging.DEBUG) 278 | stdout_handler = logging.StreamHandler(sys.stdout) 279 | logger.addHandler(stdout_handler) 280 | 281 | # Setup SQS Queue 282 | conn = boto3.client('sqs', region_name='us-east-1') 283 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 284 | 285 | # Add Slow tasks 286 | message = json.dumps({ 287 | 'task': 'tests.tasks.sleeper', 288 | 'args': [], 289 | 'kwargs': { 290 | 'message': 5, 291 | }, 292 | }) 293 | 294 | # Fill the queue (we need a lot of messages to trigger the bug) 295 | for _ in range(20): 296 | conn.send_message(QueueUrl=queue_url, MessageBody=message) 297 | 298 | # Create function to watch and kill stuck processes 299 | def sleep_and_kill(pid): 300 | import os 301 | import signal 302 | import time 303 | # This sleep time is long enoug for 100 messages in queue 304 | time.sleep(5) 305 | try: 306 | os.kill(pid, signal.SIGKILL) 307 | except OSError: 308 | # Return that we didn't need to kill the process 309 | return True 310 | else: 311 | # Return that we needed to kill the process 312 | return False 313 | 314 | # Setup Manager 315 | manager = SimpleManagerWorker( 316 | queue_prefixes=["tester"], worker_concurrency=1, interval=0.0, 317 | batchsize=1, 318 | ) 319 | manager.start() 320 | 321 | # Give our processes a moment to start 322 | time.sleep(1) 323 | 324 | # Setup Threading watcher 325 | try: 326 | # Try Python 2 Style 327 | thread = ThreadWithReturnValue2( 328 | target=sleep_and_kill, args=(manager.worker_children[0].pid,)) 329 | thread.daemon = True 330 | except TypeError: 331 | # Use Python 3 Style 332 | thread = ThreadWithReturnValue3( 333 | target=sleep_and_kill, args=(manager.worker_children[0].pid,), 334 | daemon=True, 335 | ) 336 | 337 | thread.start() 338 | 339 | # Stop the Master Process 340 | manager.stop() 341 | 342 | # Check if we had to kill the Process Worker or it exited gracefully 343 | return_value = thread.join() 344 | if not return_value: 345 | raise Exception("Process Worker failed to quit!") 346 | 347 | 348 | @mock_sqs 349 | def test_manager_picks_up_new_queues(): 350 | """ 351 | Test that the simple manager will recognize new SQS queues have been added 352 | """ 353 | 354 | # Setup SQS Queue 355 | conn = boto3.client('sqs', region_name='us-east-1') 356 | 357 | # Setup Manager 358 | manager = SimpleManagerWorker( 359 | queue_prefixes=["tester"], worker_concurrency=1, interval=1, 360 | batchsize=10, 361 | ) 362 | manager.start() 363 | 364 | # No queues found 365 | len(manager.worker_children).should.equal(0) 366 | 367 | # Create the queue 368 | conn.create_queue(QueueName="tester") 369 | manager.check_for_new_queues() 370 | 371 | # The manager should have seen the new queue was created and add a reader 372 | len(manager.worker_children).should.equal(1) 373 | manager.worker_children[0].queue_url.should.equal( 374 | "https://queue.amazonaws.com/123456789012/tester") 375 | 376 | # Cleanup 377 | manager.stop() 378 | -------------------------------------------------------------------------------- /tests/test_simple_worker.py: -------------------------------------------------------------------------------- 1 | import json 2 | import logging 3 | import time 4 | 5 | import boto3 6 | from botocore.exceptions import ClientError 7 | from moto import mock_sqs 8 | from mock import patch, Mock 9 | from pyqs.worker import ( 10 | SimpleManagerWorker, BaseProcessWorker, SimpleProcessWorker, 11 | MESSAGE_DOWNLOAD_BATCH_SIZE 12 | ) 13 | from pyqs.utils import decode_message 14 | from pyqs.events import register_event 15 | from tests.utils import MockLoggingHandler, clear_events_registry 16 | 17 | BATCHSIZE = 10 18 | INTERVAL = 0.1 19 | 20 | 21 | def _create_packed_message(task_name): 22 | # Setup SQS Queue 23 | conn = boto3.client('sqs', region_name='us-east-1') 24 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 25 | 26 | # Build the SQS message 27 | message = { 28 | 'Body': json.dumps({ 29 | 'task': task_name, 30 | 'args': [], 31 | 'kwargs': { 32 | 'message': 'Test message', 33 | }, 34 | }), 35 | "ReceiptHandle": "receipt-1234", 36 | "MessageId": "message-id-1", 37 | } 38 | 39 | packed_message = { 40 | "queue": queue_url, 41 | "message": message, 42 | "start_time": time.time(), 43 | "timeout": 30, 44 | } 45 | 46 | return queue_url, packed_message 47 | 48 | 49 | def _add_messages_to_sqs(task_name, num): 50 | # Setup SQS Queue 51 | conn = boto3.client('sqs', region_name='us-east-1') 52 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 53 | 54 | # Build the SQS message 55 | message = json.dumps({ 56 | 'task': task_name, 57 | 'args': [], 58 | 'kwargs': { 59 | 'message': 'Test message', 60 | }, 61 | }) 62 | 63 | for i in range(num): 64 | conn.send_message(QueueUrl=queue_url, MessageBody=message) 65 | 66 | return queue_url 67 | 68 | 69 | @mock_sqs 70 | def test_worker_reads_messages_from_sqs(): 71 | """ 72 | Test simple worker reads from sqs queue 73 | """ 74 | queue_url = _add_messages_to_sqs('tests.tasks.index_incrementer', 1) 75 | 76 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 77 | messages = worker.read_message() 78 | 79 | found_message_body = decode_message(messages[0]) 80 | found_message_body.should.equal({ 81 | 'task': 'tests.tasks.index_incrementer', 82 | 'args': [], 83 | 'kwargs': { 84 | 'message': 'Test message', 85 | }, 86 | }) 87 | 88 | 89 | @mock_sqs 90 | def test_worker_throws_error_when_exceeding_max_number_of_messages_for_read(): 91 | """ 92 | Test simple worker reads from sqs queue and throws error when batchsize 93 | greater than 10 94 | """ 95 | queue_url = _add_messages_to_sqs('tests.tasks.index_incrementer', 1) 96 | 97 | worker = SimpleProcessWorker(queue_url, INTERVAL, 20, parent_id=1) 98 | 99 | error_msg = "Value 20 for parameter MaxNumberOfMessages is invalid" 100 | 101 | try: 102 | worker.read_message() 103 | except ClientError as exc: 104 | str(exc).should.contain(error_msg) 105 | 106 | 107 | @mock_sqs 108 | def test_worker_reads_max_messages_from_sqs(): 109 | """ 110 | Test simple worker reads at maximum 10 message from sqs queue 111 | """ 112 | _add_messages_to_sqs('tests.tasks.index_incrementer', 12) 113 | 114 | manager = SimpleManagerWorker( 115 | queue_prefixes=['tester'], worker_concurrency=1, interval=INTERVAL, 116 | batchsize=20, 117 | ) 118 | worker = manager.worker_children[0] 119 | messages = worker.read_message() 120 | 121 | messages.should.have.length_of(BATCHSIZE) 122 | 123 | 124 | @mock_sqs 125 | def test_worker_fills_internal_queue_from_celery_task(): 126 | """ 127 | Test simple worker reads from sqs queue with celery tasks 128 | """ 129 | conn = boto3.client('sqs', region_name='us-east-1') 130 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 131 | 132 | message = ( 133 | '{"body": "KGRwMApTJ3Rhc2snCnAxClMndGVzdHMudGFza3MuaW5kZXhfa' 134 | 'W5jcmVtZW50ZXInCnAyCnNTJ2Fy\\nZ3MnCnAzCihscDQKc1Mna3dhcmdzJw' 135 | 'pwNQooZHA2ClMnbWVzc2FnZScKcDcKUydUZXN0IG1lc3Nh\\nZ2UyJwpwOAp' 136 | 'zcy4=\\n", "some stuff": "asdfasf"}' 137 | ) 138 | conn.send_message(QueueUrl=queue_url, MessageBody=message) 139 | 140 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 141 | messages = worker.read_message() 142 | 143 | found_message_body = decode_message(messages[0]) 144 | found_message_body.should.equal({ 145 | 'task': 'tests.tasks.index_incrementer', 146 | 'args': [], 147 | 'kwargs': { 148 | 'message': 'Test message2', 149 | }, 150 | }) 151 | 152 | 153 | @mock_sqs 154 | def test_worker_processes_tasks_and_logs_correctly(): 155 | """ 156 | Test simple worker processes logs INFO correctly 157 | """ 158 | # Setup logging 159 | logger = logging.getLogger("pyqs") 160 | del logger.handlers[:] 161 | logger.handlers.append(MockLoggingHandler()) 162 | 163 | queue_url, packed_message = _create_packed_message( 164 | 'tests.tasks.index_incrementer' 165 | ) 166 | 167 | # Process message 168 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 169 | worker.process_message(packed_message) 170 | 171 | # Check output 172 | kwargs = json.loads(packed_message["message"]['Body'])['kwargs'] 173 | expected_result = ( 174 | u"Processed task tests.tasks.index_incrementer in 0.0000 seconds " 175 | "with args: [] and kwargs: {}".format(kwargs) 176 | ) 177 | logger.handlers[0].messages['info'].should.equal([expected_result]) 178 | 179 | 180 | @mock_sqs 181 | def test_worker_processes_tasks_and_logs_warning_correctly(): 182 | """ 183 | Test simple worker processes logs WARNING correctly 184 | """ 185 | # Setup logging 186 | logger = logging.getLogger("pyqs") 187 | del logger.handlers[:] 188 | logger.handlers.append(MockLoggingHandler()) 189 | 190 | # Setup SQS Queue 191 | conn = boto3.client('sqs', region_name='us-east-1') 192 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 193 | 194 | # Build the SQS Message 195 | message = { 196 | 'Body': json.dumps({ 197 | 'task': 'tests.tasks.index_incrementer', 198 | 'args': [], 199 | 'kwargs': { 200 | 'message': 23, 201 | }, 202 | }), 203 | "ReceiptHandle": "receipt-1234", 204 | "MessageId": "message-id-1", 205 | } 206 | 207 | packed_message = { 208 | "queue": queue_url, 209 | "message": message, 210 | "start_time": time.time(), 211 | "timeout": 30, 212 | } 213 | 214 | # Process message 215 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 216 | worker.process_message(packed_message) 217 | 218 | # Check output 219 | kwargs = json.loads(message['Body'])['kwargs'] 220 | msg1 = ( 221 | "Task tests.tasks.index_incrementer raised error in 0.0000 seconds: " 222 | "with args: [] and kwargs: {}: " 223 | "Traceback (most recent call last)".format(kwargs) 224 | ) # noqa 225 | logger.handlers[0].messages['error'][0].lower().should.contain( 226 | msg1.lower()) 227 | msg2 = ( 228 | 'ValueError: Need to be given basestring, ' 229 | 'was given 23' 230 | ) # noqa 231 | logger.handlers[0].messages['error'][0].lower().should.contain( 232 | msg2.lower()) 233 | 234 | 235 | @patch("pyqs.worker.os") 236 | def test_parent_process_death(os): 237 | """ 238 | Test simple worker processes recognize parent process death 239 | """ 240 | os.getppid.return_value = 123 241 | 242 | worker = BaseProcessWorker(parent_id=1) 243 | worker.parent_is_alive().should.be.false 244 | 245 | 246 | @patch("pyqs.worker.os") 247 | def test_parent_process_alive(os): 248 | """ 249 | Test simple worker processes recognize when parent process is alive 250 | """ 251 | os.getppid.return_value = 1234 252 | 253 | worker = BaseProcessWorker(parent_id=1234) 254 | worker.parent_is_alive().should.be.true 255 | 256 | 257 | @mock_sqs 258 | @patch("pyqs.worker.os") 259 | def test_read_message_with_parent_process_alive_and_should_not_exit(os): 260 | """ 261 | Test simple worker processes do not exit when parent is alive and shutdown 262 | is not set when reading message 263 | """ 264 | # Setup SQS Queue 265 | conn = boto3.client('sqs', region_name='us-east-1') 266 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 267 | 268 | # Setup PPID 269 | os.getppid.return_value = 1 270 | 271 | # Setup dummy read_message 272 | def read_message(): 273 | raise Exception("Called") 274 | 275 | # When I have a parent process, and shutdown is not set 276 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 277 | worker.read_message = read_message 278 | 279 | # Then read_message() is reached 280 | worker.run.when.called_with().should.throw(Exception, "Called") 281 | 282 | 283 | @mock_sqs 284 | @patch("pyqs.worker.os") 285 | def test_read_message_with_parent_process_alive_and_should_exit(os): 286 | """ 287 | Test simple worker processes exit when parent is alive and shutdown is set 288 | when reading message 289 | """ 290 | # Setup SQS Queue 291 | conn = boto3.client('sqs', region_name='us-east-1') 292 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 293 | 294 | # Setup PPID 295 | os.getppid.return_value = 1234 296 | 297 | # When I have a parent process, and shutdown is set 298 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 299 | worker.read_message = Mock() 300 | worker.shutdown() 301 | 302 | # Then I return from run() 303 | worker.run().should.be.none 304 | 305 | 306 | @mock_sqs 307 | @patch("pyqs.worker.os") 308 | def test_read_message_with_parent_process_dead_and_should_not_exit(os): 309 | """ 310 | Test simple worker processes exit when parent is dead and shutdown is not 311 | set when reading messages 312 | """ 313 | # Setup SQS Queue 314 | conn = boto3.client('sqs', region_name='us-east-1') 315 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 316 | 317 | # Setup PPID 318 | os.getppid.return_value = 123 319 | 320 | # When I have no parent process, and shutdown is not set 321 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 322 | worker.read_message = Mock() 323 | 324 | # Then I return from run() 325 | worker.run().should.be.none 326 | 327 | 328 | @mock_sqs 329 | @patch("pyqs.worker.os") 330 | def test_process_message_with_parent_process_alive_and_should_not_exit(os): 331 | """ 332 | Test simple worker processes do not exit when parent is alive and shutdown 333 | is not set when processing message 334 | """ 335 | queue_url = _add_messages_to_sqs('tests.tasks.index_incrementer', 1) 336 | 337 | # Setup PPID 338 | os.getppid.return_value = 1 339 | 340 | # Setup dummy read_message 341 | def process_message(packed_message): 342 | raise Exception("Called") 343 | 344 | # When I have a parent process, and shutdown is not set 345 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 346 | worker.process_message = process_message 347 | 348 | # Then process_message() is reached 349 | worker.run.when.called_with().should.throw(Exception, "Called") 350 | 351 | 352 | @mock_sqs 353 | @patch("pyqs.worker.os") 354 | def test_process_message_with_parent_process_dead_and_should_not_exit(os): 355 | """ 356 | Test simple worker processes exit when parent is dead and shutdown is not 357 | set when processing message 358 | """ 359 | 360 | queue_url = _add_messages_to_sqs('tests.tasks.index_incrementer', 1) 361 | 362 | # Setup PPID 363 | os.getppid.return_value = 123 364 | 365 | # When I have a parent process, and shutdown is not set 366 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 367 | worker.process_message = Mock() 368 | 369 | # Then process_message() is reached 370 | worker.run().should.be.none 371 | 372 | 373 | @mock_sqs 374 | @patch("pyqs.worker.os") 375 | def test_process_message_with_parent_process_alive_and_should_exit(os): 376 | """ 377 | Test simple worker processes exit when parent is alive and shutdown is set 378 | set when processing message 379 | """ 380 | 381 | queue_url = _add_messages_to_sqs('tests.tasks.index_incrementer', 1) 382 | 383 | # Setup PPID 384 | os.getppid.return_value = 1 385 | 386 | # When I have a parent process, and shutdown is not set 387 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 388 | worker.process_message = Mock() 389 | worker.shutdown() 390 | 391 | # Then process_message() is reached 392 | worker.run().should.be.none 393 | 394 | 395 | @mock_sqs 396 | @patch("pyqs.worker.os") 397 | def test_worker_processes_shuts_down_after_processing_its_max_number_of_msgs( 398 | os): 399 | """ 400 | Test simple worker processes shutdown after processing maximum number 401 | of messages 402 | """ 403 | os.getppid.return_value = 1 404 | 405 | queue_url = _add_messages_to_sqs('tests.tasks.index_incrementer', 2) 406 | 407 | # When I Process messages 408 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 409 | worker._messages_to_process_before_shutdown = 2 410 | 411 | # Then I return from run() 412 | worker.run().should.be.none 413 | 414 | 415 | @mock_sqs 416 | def test_worker_negative_batch_size(): 417 | """ 418 | Test simple workers with negative batch sizes 419 | """ 420 | BATCHSIZE = -1 421 | CONCURRENCY = 1 422 | QUEUE_PREFIX = "tester" 423 | INTERVAL = 0.0 424 | conn = boto3.client('sqs', region_name='us-east-1') 425 | conn.create_queue(QueueName="tester")['QueueUrl'] 426 | 427 | worker = SimpleManagerWorker( 428 | QUEUE_PREFIX, 429 | CONCURRENCY, 430 | INTERVAL, 431 | BATCHSIZE 432 | ) 433 | worker.batchsize.should.equal(1) 434 | 435 | 436 | @mock_sqs 437 | def test_worker_to_large_batch_size(): 438 | """ 439 | Test simple workers with too large of a batch size 440 | """ 441 | BATCHSIZE = 10000 442 | CONCURRENCY = 1 443 | QUEUE_PREFIX = "tester" 444 | INTERVAL = 0.0 445 | conn = boto3.client('sqs', region_name='us-east-1') 446 | conn.create_queue(QueueName="tester")['QueueUrl'] 447 | 448 | worker = SimpleManagerWorker( 449 | QUEUE_PREFIX, 450 | CONCURRENCY, 451 | INTERVAL, 452 | BATCHSIZE 453 | ) 454 | worker.batchsize.should.equal(MESSAGE_DOWNLOAD_BATCH_SIZE) 455 | 456 | 457 | @clear_events_registry 458 | @mock_sqs 459 | def test_worker_processes_tasks_with_pre_process_callback(): 460 | """ 461 | Test simple worker runs registered callbacks when processing a message 462 | """ 463 | 464 | queue_url, packed_message = _create_packed_message( 465 | 'tests.tasks.index_incrementer' 466 | ) 467 | 468 | # Declare this so it can be checked as a side effect 469 | # to pre_process_with_side_effect 470 | contexts = [] 471 | 472 | def pre_process_with_side_effect(context): 473 | contexts.append(context) 474 | 475 | # When we have a registered pre_process callback 476 | register_event("pre_process", pre_process_with_side_effect) 477 | 478 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 479 | worker.process_message(packed_message) 480 | 481 | pre_process_context = contexts[0] 482 | 483 | # We should run the callback with the task context 484 | pre_process_context['task_name'].should.equal('index_incrementer') 485 | pre_process_context['args'].should.equal([]) 486 | pre_process_context['kwargs'].should.equal({'message': 'Test message'}) 487 | pre_process_context['full_task_path'].should.equal( 488 | 'tests.tasks.index_incrementer' 489 | ) 490 | pre_process_context['queue_url'].should.equal( 491 | 'https://queue.amazonaws.com/123456789012/tester' 492 | ) 493 | pre_process_context['timeout'].should.equal(30) 494 | 495 | assert 'fetch_time' in pre_process_context 496 | assert 'receipt_handle' in pre_process_context 497 | assert 'status' not in pre_process_context 498 | 499 | 500 | @clear_events_registry 501 | @mock_sqs 502 | def test_worker_processes_tasks_with_post_process_callback_success(): 503 | """ 504 | Test simple worker runs registered callbacks when 505 | processing a message and it succeeds 506 | """ 507 | 508 | queue_url, packed_message = _create_packed_message( 509 | 'tests.tasks.index_incrementer' 510 | ) 511 | 512 | # Declare this so it can be checked as a side effect 513 | # to post_process_with_side_effect 514 | contexts = [] 515 | 516 | def post_process_with_side_effect(context): 517 | contexts.append(context) 518 | 519 | # When we have a registered post_process callback 520 | register_event("post_process", post_process_with_side_effect) 521 | 522 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 523 | worker.process_message(packed_message) 524 | 525 | post_process_context = contexts[0] 526 | 527 | # We should run the callback with the task context 528 | post_process_context['task_name'].should.equal('index_incrementer') 529 | post_process_context['args'].should.equal([]) 530 | post_process_context['kwargs'].should.equal({'message': 'Test message'}) 531 | post_process_context['full_task_path'].should.equal( 532 | 'tests.tasks.index_incrementer' 533 | ) 534 | post_process_context['queue_url'].should.equal( 535 | 'https://queue.amazonaws.com/123456789012/tester' 536 | ) 537 | post_process_context['timeout'].should.equal(30) 538 | post_process_context['status'].should.equal('success') 539 | 540 | assert 'fetch_time' in post_process_context 541 | assert 'receipt_handle' in post_process_context 542 | assert 'exception' not in post_process_context 543 | 544 | 545 | @clear_events_registry 546 | @mock_sqs 547 | def test_worker_processes_tasks_with_post_process_callback_exception(): 548 | """ 549 | Test simple worker runs registered callbacks when processing 550 | a message and it fails 551 | """ 552 | 553 | queue_url, packed_message = _create_packed_message( 554 | 'tests.tasks.exception_task' 555 | ) 556 | 557 | # Declare this so it can be checked as a side effect 558 | # to post_process_with_side_effect 559 | contexts = [] 560 | 561 | def post_process_with_side_effect(context): 562 | contexts.append(context) 563 | 564 | # When we have a registered post_process callback 565 | register_event("post_process", post_process_with_side_effect) 566 | 567 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 568 | worker.process_message(packed_message) 569 | 570 | post_process_context = contexts[0] 571 | 572 | # We should run the callback with the task context 573 | post_process_context['task_name'].should.equal('exception_task') 574 | post_process_context['args'].should.equal([]) 575 | post_process_context['kwargs'].should.equal({'message': 'Test message'}) 576 | post_process_context['full_task_path'].should.equal( 577 | 'tests.tasks.exception_task' 578 | ) 579 | post_process_context['queue_url'].should.equal( 580 | 'https://queue.amazonaws.com/123456789012/tester' 581 | ) 582 | post_process_context['timeout'].should.equal(30) 583 | post_process_context['status'].should.equal('exception') 584 | 585 | assert 'fetch_time' in post_process_context 586 | assert 'receipt_handle' in post_process_context 587 | assert 'exception' in post_process_context 588 | 589 | 590 | @clear_events_registry 591 | @mock_sqs 592 | def test_worker_processes_tasks_with_pre_and_post_process(): 593 | """ 594 | Test worker runs registered callbacks when processing a message 595 | """ 596 | 597 | queue_url, packed_message = _create_packed_message( 598 | 'tests.tasks.index_incrementer' 599 | ) 600 | 601 | # Declare these so they can be checked as a side effect to the callbacks 602 | contexts = [] 603 | 604 | def pre_process_with_side_effect(context): 605 | contexts.append(context) 606 | 607 | def post_process_with_side_effect(context): 608 | contexts.append(context) 609 | 610 | # When we have a registered pre_process and post_process callback 611 | register_event("pre_process", pre_process_with_side_effect) 612 | register_event("post_process", post_process_with_side_effect) 613 | 614 | worker = SimpleProcessWorker(queue_url, INTERVAL, BATCHSIZE, parent_id=1) 615 | worker.process_message(packed_message) 616 | 617 | pre_process_context = contexts[0] 618 | 619 | # We should run the callbacks with the right task contexts 620 | pre_process_context['task_name'].should.equal('index_incrementer') 621 | pre_process_context['args'].should.equal([]) 622 | pre_process_context['kwargs'].should.equal({'message': 'Test message'}) 623 | pre_process_context['full_task_path'].should.equal( 624 | 'tests.tasks.index_incrementer' 625 | ) 626 | pre_process_context['queue_url'].should.equal( 627 | 'https://queue.amazonaws.com/123456789012/tester' 628 | ) 629 | pre_process_context['timeout'].should.equal(30) 630 | 631 | assert 'fetch_time' in pre_process_context 632 | assert 'receipt_handle' in pre_process_context 633 | assert 'status' not in pre_process_context 634 | 635 | post_process_context = contexts[1] 636 | 637 | post_process_context['task_name'].should.equal('index_incrementer') 638 | post_process_context['args'].should.equal([]) 639 | post_process_context['kwargs'].should.equal({'message': 'Test message'}) 640 | post_process_context['full_task_path'].should.equal( 641 | 'tests.tasks.index_incrementer' 642 | ) 643 | post_process_context['queue_url'].should.equal( 644 | 'https://queue.amazonaws.com/123456789012/tester' 645 | ) 646 | post_process_context['timeout'].should.equal(30) 647 | post_process_context['status'].should.equal('success') 648 | 649 | assert 'fetch_time' in post_process_context 650 | assert 'receipt_handle' in post_process_context 651 | assert 'exception' not in post_process_context 652 | -------------------------------------------------------------------------------- /tests/test_tasks.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | import boto3 4 | from moto import mock_sqs 5 | 6 | from .tasks import ( 7 | index_incrementer, send_email, delayed_task, custom_path_task, 8 | ) 9 | 10 | 11 | @mock_sqs() 12 | def test_basic_delay(): 13 | """ 14 | Test delaying task to default queue 15 | """ 16 | conn = boto3.client('sqs', region_name='us-east-1') 17 | conn.create_queue(QueueName="tests.tasks.index_incrementer") 18 | 19 | index_incrementer.delay("foobar", **{'extra': 'more'}) 20 | 21 | all_queues = conn.list_queues().get('QueueUrls', []) 22 | len(all_queues).should.equal(1) 23 | 24 | queue_url = all_queues[0] 25 | queue_url.should.equal( 26 | "https://queue.amazonaws.com/123456789012/" 27 | "tests.tasks.index_incrementer" 28 | ) 29 | queue = conn.get_queue_attributes( 30 | QueueUrl=queue_url, AttributeNames=['All'])['Attributes'] 31 | queue['ApproximateNumberOfMessages'].should.equal('1') 32 | 33 | message = conn.receive_message(QueueUrl=queue_url)['Messages'][0] 34 | message_dict = json.loads(message['Body']) 35 | message_dict.should.equal({ 36 | 'task': 'tests.tasks.index_incrementer', 37 | 'args': ["foobar"], 38 | 'kwargs': {'extra': 'more'}, 39 | }) 40 | 41 | 42 | @mock_sqs() 43 | def test_specified_queue(): 44 | """ 45 | Test delaying task to specific queue 46 | """ 47 | conn = boto3.client('sqs', region_name='us-east-1') 48 | 49 | send_email.delay("email subject") 50 | 51 | queue_urls = conn.list_queues().get('QueueUrls', []) 52 | len(queue_urls).should.equal(1) 53 | 54 | queue_url = queue_urls[0] 55 | queue_url.should.equal("https://queue.amazonaws.com/123456789012/email") 56 | queue = conn.get_queue_attributes( 57 | QueueUrl=queue_url, AttributeNames=['All'])['Attributes'] 58 | queue['ApproximateNumberOfMessages'].should.equal('1') 59 | 60 | 61 | @mock_sqs() 62 | def test_message_delay(): 63 | """ 64 | Test delaying task with delay_seconds 65 | """ 66 | conn = boto3.client('sqs', region_name='us-east-1') 67 | 68 | delayed_task.delay() 69 | 70 | queue_urls = conn.list_queues().get('QueueUrls', []) 71 | len(queue_urls).should.equal(1) 72 | 73 | queue_url = queue_urls[0] 74 | queue_url.should.equal("https://queue.amazonaws.com/123456789012/delayed") 75 | queue = conn.get_queue_attributes( 76 | QueueUrl=queue_url, AttributeNames=['All'])['Attributes'] 77 | queue['ApproximateNumberOfMessages'].should.equal('0') 78 | 79 | 80 | @mock_sqs() 81 | def test_message_add_delay(): 82 | """ 83 | Test configuring the delay time of a task 84 | """ 85 | conn = boto3.client('sqs', region_name='us-east-1') 86 | 87 | send_email.delay("email subject", _delay_seconds=5) 88 | 89 | queue_urls = conn.list_queues().get('QueueUrls', []) 90 | len(queue_urls).should.equal(1) 91 | 92 | queue_url = queue_urls[0] 93 | queue_url.should.equal("https://queue.amazonaws.com/123456789012/email") 94 | queue = conn.get_queue_attributes( 95 | QueueUrl=queue_url, AttributeNames=['All'])['Attributes'] 96 | queue['ApproximateNumberOfMessages'].should.equal('0') 97 | 98 | 99 | @mock_sqs() 100 | def test_message_no_delay(): 101 | """ 102 | Test removing the delay time of a task 103 | """ 104 | conn = boto3.client('sqs', region_name='us-east-1') 105 | 106 | delayed_task.delay(_delay_seconds=0) 107 | 108 | queue_urls = conn.list_queues().get('QueueUrls', []) 109 | len(queue_urls).should.equal(1) 110 | 111 | queue_url = queue_urls[0] 112 | queue_url.should.equal("https://queue.amazonaws.com/123456789012/delayed") 113 | queue = conn.get_queue_attributes( 114 | QueueUrl=queue_url, AttributeNames=['All'])['Attributes'] 115 | queue['ApproximateNumberOfMessages'].should.equal('1') 116 | 117 | 118 | @mock_sqs() 119 | def test_custom_function_path(): 120 | """ 121 | Test delaying task with custom function path 122 | """ 123 | conn = boto3.client('sqs', region_name='us-east-1') 124 | 125 | custom_path_task.delay() 126 | 127 | queue_urls = conn.list_queues().get('QueueUrls', []) 128 | len(queue_urls).should.equal(1) 129 | queue_url = queue_urls[0] 130 | queue_url.should.equal("https://queue.amazonaws.com/123456789012/foobar") 131 | queue = conn.get_queue_attributes( 132 | QueueUrl=queue_url, AttributeNames=['All'])['Attributes'] 133 | queue['ApproximateNumberOfMessages'].should.equal('1') 134 | 135 | message = conn.receive_message(QueueUrl=queue_url)['Messages'][0] 136 | message_dict = json.loads(message['Body']) 137 | message_dict.should.equal({ 138 | 'task': 'custom_function.path', 139 | 'args': [], 140 | 'kwargs': {}, 141 | }) 142 | -------------------------------------------------------------------------------- /tests/test_worker.py: -------------------------------------------------------------------------------- 1 | import json 2 | import logging 3 | import time 4 | import threading 5 | 6 | from multiprocessing import Queue 7 | try: 8 | from queue import Empty 9 | except ImportError: 10 | from Queue import Empty 11 | 12 | import boto3 13 | from moto import mock_sqs 14 | from mock import patch, Mock 15 | from pyqs.worker import ( 16 | ManagerWorker, ReadWorker, ProcessWorker, BaseWorker, 17 | MESSAGE_DOWNLOAD_BATCH_SIZE, 18 | ) 19 | from pyqs.utils import decode_message 20 | from pyqs.events import register_event 21 | from tests.tasks import task_results 22 | from tests.utils import MockLoggingHandler, clear_events_registry 23 | 24 | BATCHSIZE = 10 25 | INTERVAL = 0.1 26 | 27 | 28 | def _add_message_to_internal_queue(task_name): 29 | # Setup SQS Queue 30 | conn = boto3.client('sqs', region_name='us-east-1') 31 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 32 | 33 | # Build the SQS message 34 | message = { 35 | 'Body': json.dumps({ 36 | 'task': task_name, 37 | 'args': [], 38 | 'kwargs': { 39 | 'message': 'Test message', 40 | }, 41 | }), 42 | "ReceiptHandle": "receipt-1234", 43 | "MessageId": "message-id-1", 44 | } 45 | # Add message to queue 46 | internal_queue = Queue() 47 | internal_queue.put( 48 | { 49 | "message": message, 50 | "queue": queue_url, 51 | "start_time": time.time(), 52 | "timeout": 30, 53 | } 54 | ) 55 | return internal_queue 56 | 57 | 58 | def _check_internal_queue_is_empty(internal_queue): 59 | try: 60 | internal_queue.get(timeout=1) 61 | except Empty: 62 | pass 63 | else: 64 | raise AssertionError("The internal queue should be empty") 65 | 66 | 67 | @mock_sqs 68 | def test_worker_fills_internal_queue(): 69 | """ 70 | Test read workers fill internal queue 71 | """ 72 | conn = boto3.client('sqs', region_name='us-east-1') 73 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 74 | 75 | message = json.dumps({ 76 | 'task': 'tests.tasks.index_incrementer', 77 | 'args': [], 78 | 'kwargs': { 79 | 'message': 'Test message', 80 | }, 81 | }) 82 | 83 | conn.send_message(QueueUrl=queue_url, MessageBody=message) 84 | 85 | internal_queue = Queue() 86 | worker = ReadWorker(queue_url, internal_queue, BATCHSIZE, parent_id=1) 87 | worker.read_message() 88 | 89 | packed_message = internal_queue.get(timeout=1) 90 | found_message_body = decode_message(packed_message['message']) 91 | found_message_body.should.equal({ 92 | 'task': 'tests.tasks.index_incrementer', 93 | 'args': [], 94 | 'kwargs': { 95 | 'message': 'Test message', 96 | }, 97 | }) 98 | 99 | 100 | @mock_sqs 101 | def test_worker_fills_internal_queue_only_until_maximum_queue_size(): 102 | """ 103 | Test read workers fill internal queue only to maximum size 104 | """ 105 | conn = boto3.client('sqs', region_name='us-east-1') 106 | # Set visibility timeout low to improve test speed 107 | queue_url = conn.create_queue( 108 | QueueName="tester", Attributes={'VisibilityTimeout': '1'})['QueueUrl'] 109 | 110 | message = json.dumps({ 111 | 'task': 'tests.tasks.index_incrementer', 112 | 'args': [], 113 | 'kwargs': { 114 | 'message': 'Test message', 115 | }, 116 | }) 117 | for i in range(3): 118 | conn.send_message(QueueUrl=queue_url, MessageBody=message) 119 | 120 | internal_queue = Queue(maxsize=2) 121 | worker = ReadWorker(queue_url, internal_queue, BATCHSIZE, parent_id=1) 122 | worker.read_message() 123 | 124 | # The internal queue should only have two messages on it 125 | internal_queue.get(timeout=1) 126 | internal_queue.get(timeout=1) 127 | 128 | try: 129 | internal_queue.get(timeout=1) 130 | except Empty: 131 | pass 132 | else: 133 | raise AssertionError("The internal queue should be empty") 134 | 135 | 136 | @mock_sqs 137 | def test_worker_fills_internal_queue_from_celery_task(): 138 | """ 139 | Test read workers fill internal queue with celery tasks 140 | """ 141 | conn = boto3.client('sqs', region_name='us-east-1') 142 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 143 | 144 | message = ( 145 | '{"body": "KGRwMApTJ3Rhc2snCnAxClMndGVzdHMudGFza3MuaW5kZXhfa' 146 | 'W5jcmVtZW50ZXInCnAyCnNTJ2Fy\\nZ3MnCnAzCihscDQKc1Mna3dhcmdzJw' 147 | 'pwNQooZHA2ClMnbWVzc2FnZScKcDcKUydUZXN0IG1lc3Nh\\nZ2UyJwpwOAp' 148 | 'zcy4=\\n", "some stuff": "asdfasf"}' 149 | ) 150 | conn.send_message(QueueUrl=queue_url, MessageBody=message) 151 | 152 | internal_queue = Queue() 153 | worker = ReadWorker(queue_url, internal_queue, BATCHSIZE, parent_id=1) 154 | worker.read_message() 155 | 156 | packed_message = internal_queue.get(timeout=1) 157 | found_message_body = decode_message(packed_message['message']) 158 | found_message_body.should.equal({ 159 | 'task': 'tests.tasks.index_incrementer', 160 | 'args': [], 161 | 'kwargs': { 162 | 'message': 'Test message2', 163 | }, 164 | }) 165 | 166 | 167 | @mock_sqs 168 | def test_worker_processes_tasks_from_internal_queue(): 169 | """ 170 | Test worker processes read from internal queue 171 | """ 172 | del task_results[:] 173 | 174 | # Setup SQS Queue 175 | conn = boto3.client('sqs', region_name='us-east-1') 176 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 177 | 178 | # Build the SQS message 179 | message = { 180 | 'Body': json.dumps({ 181 | 'task': 'tests.tasks.index_incrementer', 182 | 'args': [], 183 | 'kwargs': { 184 | 'message': 'Test message', 185 | }, 186 | }), 187 | "ReceiptHandle": "receipt-1234", 188 | "MessageId": "message-id-1", 189 | } 190 | 191 | # Add message to queue 192 | internal_queue = Queue() 193 | internal_queue.put( 194 | { 195 | "message": message, 196 | "queue": queue_url, 197 | "start_time": time.time(), 198 | "timeout": 30, 199 | } 200 | ) 201 | 202 | # Process message 203 | worker = ProcessWorker(internal_queue, INTERVAL, parent_id=1) 204 | worker.process_message() 205 | 206 | task_results.should.equal(['Test message']) 207 | 208 | # We expect the queue to be empty now 209 | try: 210 | internal_queue.get(timeout=1) 211 | except Empty: 212 | pass 213 | else: 214 | raise AssertionError("The internal queue should be empty") 215 | 216 | 217 | @mock_sqs 218 | def test_worker_fills_internal_queue_and_respects_visibility_timeouts(): 219 | """ 220 | Test read workers respect visibility timeouts 221 | """ 222 | # Setup logging 223 | logger = logging.getLogger("pyqs") 224 | logger.handlers.append(MockLoggingHandler()) 225 | 226 | # Setup SQS Queue 227 | conn = boto3.client('sqs', region_name='us-east-1') 228 | queue_url = conn.create_queue( 229 | QueueName="tester", Attributes={'VisibilityTimeout': '1'})['QueueUrl'] 230 | 231 | # Add MEssages 232 | message = json.dumps( 233 | { 234 | "body": ( 235 | "KGRwMApTJ3Rhc2snCnAxClMndGVzdHMudGFza3MuaW5kZXhfaW5jcmVtZW" 236 | "50ZXInCnAyCnNTJ2Fy\nZ3MnCnAzCihscDQKc1Mna3dhcmdzJwpwNQooZHA" 237 | "2ClMnbWVzc2FnZScKcDcKUydUZXN0IG1lc3Nh\nZ2UyJwpwOApzcy4=\n" 238 | ), 239 | "some stuff": "asdfasf", 240 | } 241 | ) 242 | for _ in range(3): 243 | conn.send_message(QueueUrl=queue_url, MessageBody=message) 244 | 245 | # Run Reader 246 | internal_queue = Queue(maxsize=1) 247 | worker = ReadWorker(queue_url, internal_queue, BATCHSIZE, parent_id=1) 248 | worker.read_message() 249 | 250 | # Check log messages 251 | logger.handlers[0].messages['warning'][0].should.contain( 252 | "Timed out trying to add the following message to the internal queue") 253 | logger.handlers[0].messages['warning'][1].should.contain( 254 | "Clearing Local messages since we exceeded their visibility_timeout") 255 | 256 | 257 | @mock_sqs 258 | def test_worker_processes_tasks_and_logs_correctly(): 259 | """ 260 | Test worker processes logs INFO correctly 261 | """ 262 | # Setup logging 263 | logger = logging.getLogger("pyqs") 264 | del logger.handlers[:] 265 | logger.handlers.append(MockLoggingHandler()) 266 | 267 | # Setup SQS Queue 268 | conn = boto3.client('sqs', region_name='us-east-1') 269 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 270 | 271 | # Build the SQS message 272 | message = { 273 | 'Body': json.dumps({ 274 | 'task': 'tests.tasks.index_incrementer', 275 | 'args': [], 276 | 'kwargs': { 277 | 'message': 'Test message', 278 | }, 279 | }), 280 | "ReceiptHandle": "receipt-1234", 281 | "MessageId": "message-id-1", 282 | } 283 | 284 | # Add message to internal queue 285 | internal_queue = Queue() 286 | internal_queue.put( 287 | { 288 | "queue": queue_url, 289 | "message": message, 290 | "start_time": time.time(), 291 | "timeout": 30, 292 | } 293 | ) 294 | 295 | # Process message 296 | worker = ProcessWorker(internal_queue, INTERVAL, parent_id=1) 297 | worker.process_message() 298 | 299 | # Check output 300 | kwargs = json.loads(message['Body'])['kwargs'] 301 | expected_result = ( 302 | u"Processed task tests.tasks.index_incrementer in 0.0000 seconds " 303 | "with args: [] and kwargs: {}".format(kwargs) 304 | ) 305 | logger.handlers[0].messages['info'].should.equal([expected_result]) 306 | 307 | 308 | @mock_sqs 309 | def test_worker_processes_tasks_and_logs_warning_correctly(): 310 | """ 311 | Test worker processes logs WARNING correctly 312 | """ 313 | # Setup logging 314 | logger = logging.getLogger("pyqs") 315 | del logger.handlers[:] 316 | logger.handlers.append(MockLoggingHandler()) 317 | 318 | # Setup SQS Queue 319 | conn = boto3.client('sqs', region_name='us-east-1') 320 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 321 | 322 | # Build the SQS Message 323 | message = { 324 | 'Body': json.dumps({ 325 | 'task': 'tests.tasks.index_incrementer', 326 | 'args': [], 327 | 'kwargs': { 328 | 'message': 23, 329 | }, 330 | }), 331 | "ReceiptHandle": "receipt-1234", 332 | "MessageId": "message-id-1", 333 | } 334 | 335 | # Add message to internal queue 336 | internal_queue = Queue() 337 | internal_queue.put( 338 | { 339 | "queue": queue_url, 340 | "message": message, 341 | "start_time": time.time(), 342 | "timeout": 30, 343 | } 344 | ) 345 | 346 | # Process message 347 | worker = ProcessWorker(internal_queue, INTERVAL, parent_id=1) 348 | worker.process_message() 349 | 350 | # Check output 351 | kwargs = json.loads(message['Body'])['kwargs'] 352 | msg1 = ( 353 | "Task tests.tasks.index_incrementer raised error in 0.0000 seconds: " 354 | "with args: [] and kwargs: {}: " 355 | "Traceback (most recent call last)".format(kwargs) 356 | ) # noqa 357 | logger.handlers[0].messages['error'][0].lower().should.contain( 358 | msg1.lower()) 359 | msg2 = ( 360 | 'ValueError: Need to be given basestring, ' 361 | 'was given 23' 362 | ) # noqa 363 | logger.handlers[0].messages['error'][0].lower().should.contain( 364 | msg2.lower()) 365 | 366 | 367 | @mock_sqs 368 | def test_worker_processes_empty_queue(): 369 | """ 370 | Test worker processes read from empty internal queue 371 | """ 372 | internal_queue = Queue() 373 | 374 | worker = ProcessWorker(internal_queue, INTERVAL, parent_id=1) 375 | worker.process_message() 376 | 377 | 378 | @patch("pyqs.worker.os") 379 | def test_parent_process_death(os): 380 | """ 381 | Test worker processes recognize parent process death 382 | """ 383 | os.getppid.return_value = 123 384 | 385 | worker = BaseWorker(parent_id=1) 386 | worker.parent_is_alive().should.be.false 387 | 388 | 389 | @patch("pyqs.worker.os") 390 | def test_parent_process_alive(os): 391 | """ 392 | Test worker processes recognize when parent process is alive 393 | """ 394 | os.getppid.return_value = 1234 395 | 396 | worker = BaseWorker(parent_id=1234) 397 | worker.parent_is_alive().should.be.true 398 | 399 | 400 | @mock_sqs 401 | @patch("pyqs.worker.os") 402 | def test_read_worker_with_parent_process_alive_and_should_not_exit(os): 403 | """ 404 | Test read workers do not exit when parent is alive and shutdown is not set 405 | """ 406 | # Setup SQS Queue 407 | conn = boto3.client('sqs', region_name='us-east-1') 408 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 409 | 410 | # Setup PPID 411 | os.getppid.return_value = 1 412 | 413 | # Setup dummy read_message 414 | def read_message(): 415 | raise Exception("Called") 416 | 417 | # When I have a parent process, and shutdown is not set 418 | worker = ReadWorker(queue_url, "foo", BATCHSIZE, parent_id=1) 419 | worker.read_message = read_message 420 | 421 | # Then read_message() is reached 422 | worker.run.when.called_with().should.throw(Exception, "Called") 423 | 424 | 425 | @mock_sqs 426 | @patch("pyqs.worker.os") 427 | def test_read_worker_with_parent_process_alive_and_should_exit(os): 428 | """ 429 | Test read workers exit when parent is alive and shutdown is set 430 | """ 431 | # Setup SQS Queue 432 | conn = boto3.client('sqs', region_name='us-east-1') 433 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 434 | 435 | # Setup PPID 436 | os.getppid.return_value = 1234 437 | 438 | # Setup internal queue 439 | q = Queue(1) 440 | 441 | # When I have a parent process, and shutdown is set 442 | worker = ReadWorker(queue_url, q, BATCHSIZE, parent_id=1) 443 | worker.read_message = Mock() 444 | worker.shutdown() 445 | 446 | # Then I return from run() 447 | worker.run().should.be.none 448 | 449 | 450 | @mock_sqs 451 | @patch("pyqs.worker.os") 452 | def test_read_worker_with_parent_process_dead_and_should_not_exit(os): 453 | """ 454 | Test read workers exit when parent is dead and shutdown is not set 455 | """ 456 | # Setup SQS Queue 457 | conn = boto3.client('sqs', region_name='us-east-1') 458 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 459 | 460 | # Setup PPID 461 | os.getppid.return_value = 123 462 | 463 | # Setup internal queue 464 | q = Queue(1) 465 | 466 | # When I have no parent process, and shutdown is not set 467 | worker = ReadWorker(queue_url, q, BATCHSIZE, parent_id=1) 468 | worker.read_message = Mock() 469 | 470 | # Then I return from run() 471 | worker.run().should.be.none 472 | 473 | 474 | @mock_sqs 475 | @patch("pyqs.worker.os") 476 | def test_process_worker_with_parent_process_alive_and_should_not_exit(os): 477 | """ 478 | Test worker processes do not exit when parent is alive and shutdown 479 | is not set 480 | """ 481 | # Setup PPID 482 | os.getppid.return_value = 1 483 | 484 | # Setup dummy read_message 485 | def process_message(): 486 | raise Exception("Called") 487 | 488 | # When I have a parent process, and shutdown is not set 489 | worker = ProcessWorker("foo", INTERVAL, parent_id=1) 490 | worker.process_message = process_message 491 | 492 | # Then process_message() is reached 493 | worker.run.when.called_with().should.throw(Exception, "Called") 494 | 495 | 496 | @mock_sqs 497 | @patch("pyqs.worker.os") 498 | def test_process_worker_with_parent_process_dead_and_should_not_exit(os): 499 | """ 500 | Test worker processes exit when parent is dead and shutdown is not set 501 | """ 502 | # Setup PPID 503 | os.getppid.return_value = 1 504 | 505 | # When I have no parent process, and shutdown is not set 506 | worker = ProcessWorker("foo", INTERVAL, parent_id=1) 507 | worker.process_message = Mock() 508 | 509 | # Then I return from run() 510 | worker.run().should.be.none 511 | 512 | 513 | @mock_sqs 514 | @patch("pyqs.worker.os") 515 | def test_process_worker_with_parent_process_alive_and_should_exit(os): 516 | """ 517 | Test worker processes exit when parent is alive and shutdown is set 518 | """ 519 | # Setup PPID 520 | os.getppid.return_value = 1234 521 | 522 | # When I have a parent process, and shutdown is set 523 | worker = ProcessWorker("foo", INTERVAL, parent_id=1) 524 | worker.process_message = Mock() 525 | worker.shutdown() 526 | 527 | # Then I return from run() 528 | worker.run().should.be.none 529 | 530 | 531 | @mock_sqs 532 | @patch("pyqs.worker.os") 533 | def test_worker_processes_shuts_down_after_processing_its_max_number_of_msgs( 534 | os): 535 | """ 536 | Test worker processes shutdown after processing maximum number of messages 537 | """ 538 | os.getppid.return_value = 1 539 | 540 | # Setup SQS Queue 541 | conn = boto3.client('sqs', region_name='us-east-1') 542 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 543 | 544 | # Build the SQS Message 545 | message = { 546 | 'Body': json.dumps({ 547 | 'task': 'tests.tasks.index_incrementer', 548 | 'args': [], 549 | 'kwargs': { 550 | 'message': 23, 551 | }, 552 | }), 553 | "ReceiptHandle": "receipt-1234", 554 | "MessageId": "message-id-1", 555 | } 556 | 557 | # Add message to internal queue 558 | internal_queue = Queue(3) 559 | internal_queue.put( 560 | { 561 | "queue": queue_url, 562 | "message": message, 563 | "start_time": time.time(), 564 | "timeout": 30, 565 | } 566 | ) 567 | internal_queue.put( 568 | { 569 | "queue": queue_url, 570 | "message": message, 571 | "start_time": time.time(), 572 | "timeout": 30, 573 | } 574 | ) 575 | internal_queue.put( 576 | { 577 | "queue": queue_url, 578 | "message": message, 579 | "start_time": time.time(), 580 | "timeout": 30, 581 | } 582 | ) 583 | 584 | # When I Process messages 585 | worker = ProcessWorker(internal_queue, INTERVAL, parent_id=1) 586 | worker._messages_to_process_before_shutdown = 2 587 | 588 | # Then I return from run() 589 | worker.run().should.be.none 590 | 591 | # With messages still on the queue 592 | internal_queue.empty().should.be.false 593 | internal_queue.full().should.be.false 594 | 595 | 596 | @mock_sqs 597 | def test_worker_processes_discard_tasks_that_exceed_their_visibility_timeout(): 598 | """ 599 | Test worker processes discards tasks that exceed their visibility timeout 600 | """ 601 | # Setup logging 602 | logger = logging.getLogger("pyqs") 603 | del logger.handlers[:] 604 | logger.handlers.append(MockLoggingHandler()) 605 | 606 | # Setup SQS Queue 607 | conn = boto3.client('sqs', region_name='us-east-1') 608 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 609 | 610 | # Build the SQS Message 611 | message = { 612 | 'Body': json.dumps({ 613 | 'task': 'tests.tasks.index_incrementer', 614 | 'args': [], 615 | 'kwargs': { 616 | 'message': 23, 617 | }, 618 | }), 619 | "ReceiptHandle": "receipt-1234", 620 | "MessageId": "message-id-1", 621 | } 622 | 623 | # Add message to internal queue with timeout of 0 that started long ago 624 | internal_queue = Queue() 625 | internal_queue.put( 626 | { 627 | "queue": queue_url, 628 | "message": message, 629 | "start_time": 0, 630 | "timeout": 0, 631 | } 632 | ) 633 | 634 | # When I process the message 635 | worker = ProcessWorker(internal_queue, INTERVAL, parent_id=1) 636 | worker.process_message() 637 | 638 | # Then I get an error about exceeding the visibility timeout 639 | kwargs = json.loads(message['Body'])['kwargs'] 640 | msg1 = ( 641 | "Discarding task tests.tasks.index_incrementer with args: [] " 642 | "and kwargs: {} due to exceeding " 643 | "visibility timeout" 644 | ).format(kwargs) # noqa 645 | logger.handlers[0].messages['warning'][0].lower().should.contain( 646 | msg1.lower()) 647 | 648 | 649 | @mock_sqs 650 | def test_worker_processes_only_incr_processed_counter_if_a_msg_was_processed(): 651 | """ 652 | Test worker process only increases processed counter if a message was 653 | processed 654 | """ 655 | # Setup SQS Queue 656 | conn = boto3.client('sqs', region_name='us-east-1') 657 | queue_url = conn.create_queue(QueueName="tester")['QueueUrl'] 658 | 659 | # Build the SQS Message 660 | message = { 661 | 'Body': json.dumps({ 662 | 'task': 'tests.tasks.index_incrementer', 663 | 'args': [], 664 | 'kwargs': { 665 | 'message': 23, 666 | }, 667 | }), 668 | "ReceiptHandle": "receipt-1234", 669 | "MessageId": "message-id-1", 670 | } 671 | 672 | # Add message to internal queue 673 | internal_queue = Queue(3) 674 | internal_queue.put( 675 | { 676 | "queue": queue_url, 677 | "message": message, 678 | "start_time": time.time(), 679 | "timeout": 30, 680 | } 681 | ) 682 | 683 | # And we add a message to the queue later 684 | def sleep_and_queue(internal_queue): 685 | time.sleep(1) 686 | internal_queue.put( 687 | { 688 | "queue": queue_url, 689 | "message": message, 690 | "start_time": time.time(), 691 | "timeout": 30, 692 | } 693 | ) 694 | 695 | thread = threading.Thread(target=sleep_and_queue, args=(internal_queue,)) 696 | thread.daemon = True 697 | thread.start() 698 | 699 | # When I Process messages 700 | worker = ProcessWorker(internal_queue, INTERVAL, parent_id=1) 701 | worker._messages_to_process_before_shutdown = 2 702 | 703 | # Then I return from run() after processing 2 messages 704 | worker.run().should.be.none 705 | 706 | 707 | @mock_sqs 708 | def test_worker_negative_batch_size(): 709 | """ 710 | Test workers with negative batch sizes 711 | """ 712 | BATCHSIZE = -1 713 | CONCURRENCY = 1 714 | QUEUE_PREFIX = "tester" 715 | INTERVAL = 0.0 716 | conn = boto3.client('sqs', region_name='us-east-1') 717 | conn.create_queue(QueueName="tester")['QueueUrl'] 718 | 719 | worker = ManagerWorker(QUEUE_PREFIX, CONCURRENCY, INTERVAL, BATCHSIZE) 720 | worker.batchsize.should.equal(1) 721 | 722 | 723 | @mock_sqs 724 | def test_worker_to_large_batch_size(): 725 | """ 726 | Test workers with too large of a batch size 727 | """ 728 | BATCHSIZE = 10000 729 | CONCURRENCY = 1 730 | QUEUE_PREFIX = "tester" 731 | INTERVAL = 0.0 732 | conn = boto3.client('sqs', region_name='us-east-1') 733 | conn.create_queue(QueueName="tester")['QueueUrl'] 734 | 735 | worker = ManagerWorker(QUEUE_PREFIX, CONCURRENCY, INTERVAL, BATCHSIZE) 736 | worker.batchsize.should.equal(MESSAGE_DOWNLOAD_BATCH_SIZE) 737 | 738 | 739 | @clear_events_registry 740 | @mock_sqs 741 | def test_worker_processes_tasks_with_pre_process_callback(): 742 | """ 743 | Test worker runs registered callbacks when processing a message 744 | """ 745 | 746 | # Declare this so it can be checked as a side effect 747 | # to pre_process_with_side_effect 748 | contexts = [] 749 | 750 | def pre_process_with_side_effect(context): 751 | contexts.append(context) 752 | 753 | # When we have a registered pre_process callback 754 | register_event("pre_process", pre_process_with_side_effect) 755 | 756 | # And we process a message 757 | internal_queue = _add_message_to_internal_queue( 758 | 'tests.tasks.index_incrementer' 759 | ) 760 | worker = ProcessWorker(internal_queue, INTERVAL, parent_id=1) 761 | worker.process_message() 762 | 763 | pre_process_context = contexts[0] 764 | 765 | # We should run the callback with the task context 766 | pre_process_context['task_name'].should.equal('index_incrementer') 767 | pre_process_context['args'].should.equal([]) 768 | pre_process_context['kwargs'].should.equal({'message': 'Test message'}) 769 | pre_process_context['full_task_path'].should.equal( 770 | 'tests.tasks.index_incrementer' 771 | ) 772 | pre_process_context['queue_url'].should.equal( 773 | 'https://queue.amazonaws.com/123456789012/tester' 774 | ) 775 | pre_process_context['timeout'].should.equal(30) 776 | 777 | assert 'fetch_time' in pre_process_context 778 | assert 'status' not in pre_process_context 779 | 780 | # And the internal queue should be empty 781 | _check_internal_queue_is_empty(internal_queue) 782 | 783 | 784 | @clear_events_registry 785 | @mock_sqs 786 | def test_worker_processes_tasks_with_post_process_callback_success(): 787 | """ 788 | Test worker runs registered callbacks when 789 | processing a message and it succeeds 790 | """ 791 | 792 | # Declare this so it can be checked as a side effect 793 | # to post_process_with_side_effect 794 | contexts = [] 795 | 796 | def post_process_with_side_effect(context): 797 | contexts.append(context) 798 | 799 | # When we have a registered post_process callback 800 | register_event("post_process", post_process_with_side_effect) 801 | 802 | # And we process a message 803 | internal_queue = _add_message_to_internal_queue( 804 | 'tests.tasks.index_incrementer' 805 | ) 806 | worker = ProcessWorker(internal_queue, INTERVAL, parent_id=1) 807 | worker.process_message() 808 | 809 | post_process_context = contexts[0] 810 | 811 | # We should run the callback with the task context 812 | post_process_context['task_name'].should.equal('index_incrementer') 813 | post_process_context['args'].should.equal([]) 814 | post_process_context['kwargs'].should.equal({'message': 'Test message'}) 815 | post_process_context['full_task_path'].should.equal( 816 | 'tests.tasks.index_incrementer' 817 | ) 818 | post_process_context['queue_url'].should.equal( 819 | 'https://queue.amazonaws.com/123456789012/tester' 820 | ) 821 | post_process_context['timeout'].should.equal(30) 822 | post_process_context['status'].should.equal('success') 823 | 824 | assert 'fetch_time' in post_process_context 825 | assert 'exception' not in post_process_context 826 | 827 | # And the internal queue should be empty 828 | _check_internal_queue_is_empty(internal_queue) 829 | 830 | 831 | @clear_events_registry 832 | @mock_sqs 833 | def test_worker_processes_tasks_with_post_process_callback_exception(): 834 | """ 835 | Test worker runs registered callbacks when processing 836 | a message and it fails 837 | """ 838 | 839 | # Declare this so it can be checked as a side effect 840 | # to post_process_with_side_effect 841 | contexts = [] 842 | 843 | def post_process_with_side_effect(context): 844 | contexts.append(context) 845 | 846 | # When we have a registered post_process callback 847 | register_event("post_process", post_process_with_side_effect) 848 | 849 | # And we process a message 850 | internal_queue = _add_message_to_internal_queue( 851 | 'tests.tasks.exception_task' 852 | ) 853 | worker = ProcessWorker(internal_queue, INTERVAL, parent_id=1) 854 | worker.process_message() 855 | 856 | post_process_context = contexts[0] 857 | 858 | # We should run the callback with the task context 859 | post_process_context['task_name'].should.equal('exception_task') 860 | post_process_context['args'].should.equal([]) 861 | post_process_context['kwargs'].should.equal({'message': 'Test message'}) 862 | post_process_context['full_task_path'].should.equal( 863 | 'tests.tasks.exception_task' 864 | ) 865 | post_process_context['queue_url'].should.equal( 866 | 'https://queue.amazonaws.com/123456789012/tester' 867 | ) 868 | post_process_context['timeout'].should.equal(30) 869 | post_process_context['status'].should.equal('exception') 870 | 871 | assert 'fetch_time' in post_process_context 872 | assert 'exception' in post_process_context 873 | 874 | # And the internal queue should be empty 875 | _check_internal_queue_is_empty(internal_queue) 876 | 877 | 878 | @clear_events_registry 879 | @mock_sqs 880 | def test_worker_processes_tasks_with_pre_and_post_process(): 881 | """ 882 | Test worker runs registered callbacks when processing a message 883 | """ 884 | 885 | # Declare these so they can be checked as a side effect to the callbacks 886 | contexts = [] 887 | 888 | def pre_process_with_side_effect(context): 889 | contexts.append(context) 890 | 891 | def post_process_with_side_effect(context): 892 | contexts.append(context) 893 | 894 | # When we have a registered pre_process and post_process callback 895 | register_event("pre_process", pre_process_with_side_effect) 896 | register_event("post_process", post_process_with_side_effect) 897 | 898 | # And we process a message 899 | internal_queue = _add_message_to_internal_queue( 900 | 'tests.tasks.index_incrementer' 901 | ) 902 | worker = ProcessWorker(internal_queue, INTERVAL, parent_id=1) 903 | worker.process_message() 904 | 905 | pre_process_context = contexts[0] 906 | 907 | # We should run the callbacks with the right task contexts 908 | pre_process_context['task_name'].should.equal('index_incrementer') 909 | pre_process_context['args'].should.equal([]) 910 | pre_process_context['kwargs'].should.equal({'message': 'Test message'}) 911 | pre_process_context['full_task_path'].should.equal( 912 | 'tests.tasks.index_incrementer' 913 | ) 914 | pre_process_context['queue_url'].should.equal( 915 | 'https://queue.amazonaws.com/123456789012/tester' 916 | ) 917 | pre_process_context['timeout'].should.equal(30) 918 | 919 | assert 'fetch_time' in pre_process_context 920 | assert 'status' not in pre_process_context 921 | 922 | post_process_context = contexts[1] 923 | 924 | post_process_context['task_name'].should.equal('index_incrementer') 925 | post_process_context['args'].should.equal([]) 926 | post_process_context['kwargs'].should.equal({'message': 'Test message'}) 927 | post_process_context['full_task_path'].should.equal( 928 | 'tests.tasks.index_incrementer' 929 | ) 930 | post_process_context['queue_url'].should.equal( 931 | 'https://queue.amazonaws.com/123456789012/tester' 932 | ) 933 | post_process_context['timeout'].should.equal(30) 934 | post_process_context['status'].should.equal('success') 935 | 936 | assert 'fetch_time' in post_process_context 937 | assert 'exception' not in post_process_context 938 | 939 | # And the internal queue should be empty 940 | _check_internal_queue_is_empty(internal_queue) 941 | -------------------------------------------------------------------------------- /tests/utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | from __future__ import unicode_literals 3 | 4 | import logging 5 | from functools import wraps 6 | 7 | from threading import Thread 8 | 9 | from pyqs import events 10 | 11 | 12 | class MockLoggingHandler(logging.Handler): 13 | """Mock logging handler to check for expected logs.""" 14 | 15 | def __init__(self, *args, **kwargs): 16 | self.reset() 17 | logging.Handler.__init__(self, *args, **kwargs) 18 | 19 | def emit(self, record): 20 | self.messages[record.levelname.lower()].append(record.getMessage()) 21 | 22 | def reset(self): 23 | self.messages = { 24 | 'debug': [], 25 | 'info': [], 26 | 'warning': [], 27 | 'error': [], 28 | 'critical': [], 29 | } 30 | 31 | 32 | class ThreadWithReturnValue2(Thread): 33 | def __init__(self, group=None, target=None, name=None, args=(), kwargs={}, 34 | Verbose=None): 35 | Thread.__init__(self, group, target, name, args, kwargs, Verbose) 36 | self._return = None 37 | 38 | def run(self): 39 | if self._Thread__target is not None: 40 | self._return = self._Thread__target( 41 | *self._Thread__args, **self._Thread__kwargs) 42 | 43 | def join(self): 44 | Thread.join(self) 45 | return self._return 46 | 47 | 48 | class ThreadWithReturnValue3(Thread): 49 | def __init__(self, group=None, target=None, name=None, args=(), 50 | kwargs=None, daemon=None): 51 | Thread.__init__(self, group, target, name, args, kwargs, daemon=daemon) 52 | self._return = None 53 | 54 | def run(self): 55 | if self._target is not None: 56 | self._return = self._target(*self._args, **self._kwargs) 57 | 58 | def join(self): 59 | Thread.join(self) 60 | return self._return 61 | 62 | 63 | def clear_events_registry(fn): 64 | """Clear the global events registry before each test.""" 65 | 66 | @wraps(fn) 67 | def wrapper(*args, **kwargs): 68 | events.clear_events() 69 | return fn(*args, **kwargs) 70 | return wrapper 71 | --------------------------------------------------------------------------------