├── .bumpversion.cfg ├── .gitignore ├── LICENSE ├── README.md ├── dramatiq_sqs ├── __init__.py └── broker.py ├── setup.cfg ├── setup.py └── tests ├── __init__.py ├── conftest.py ├── test_broker.py └── test_utils.py /.bumpversion.cfg: -------------------------------------------------------------------------------- 1 | [bumpversion] 2 | current_version = 0.2.0 3 | message = chore: bump version {current_version} → {new_version} 4 | commit = True 5 | tag = True 6 | 7 | [bumpversion:file:dramatiq_sqs/__init__.py] 8 | search = __version__ = "{current_version}" 9 | replace = __version__ = "{new_version}" 10 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.egg-info 2 | .coverage 3 | .mypy_cache 4 | .pytest_cache 5 | __pycache__ 6 | build 7 | dist 8 | htmlcov 9 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright 2018 Bogdan Paul Popa 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | http://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # dramatiq_sqs 2 | 3 | A [Dramatiq] broker that can be used with [Amazon SQS]. 4 | 5 | This backend has a number of limitations compared to the built-in 6 | Redis and RMQ backends: 7 | 8 | * the max amount of time messages can be delayed by is 15 minutes, 9 | * messages can be at most 256KiB large and 10 | * messages must be processed within 2 hours of being pulled, otherwise 11 | they will be redelivered. 12 | 13 | The backend uses [boto3] under the hood. For details on how 14 | authorization works, check out its [docs]. 15 | 16 | 17 | ## Installation 18 | 19 | pip install dramatiq_sqs 20 | 21 | 22 | ## Usage 23 | 24 | ``` python 25 | import dramatiq 26 | 27 | from dramatiq.middleware import AgeLimit, TimeLimit, Callbacks, Pipelines, Prometheus, Retries 28 | from dramatiq_sqs import SQSBroker 29 | 30 | broker = SQSBroker( 31 | namespace="dramatiq_sqs_tests", 32 | middleware=[ 33 | Prometheus(), 34 | AgeLimit(), 35 | TimeLimit(), 36 | Callbacks(), 37 | Pipelines(), 38 | Retries(min_backoff=1000, max_backoff=900000, max_retries=96), 39 | ], 40 | ) 41 | dramatiq.set_broker(broker) 42 | ``` 43 | 44 | 45 | ## Usage with [ElasticMQ] 46 | 47 | ``` python 48 | broker = SQSBroker( 49 | # ... 50 | endpoint_url="http://127.0.0.1:9324", 51 | ) 52 | ``` 53 | 54 | ## Example IAM Policy 55 | 56 | Here are the IAM permissions needed by Dramatiq: 57 | 58 | ``` json 59 | { 60 | "Version": "2012-10-17", 61 | "Statement": [ 62 | { 63 | "Effect": "Allow", 64 | "Action": [ 65 | "sqs:CreateQueue", 66 | "sqs:ReceiveMessage", 67 | "sqs:DeleteMessage", 68 | "sqs:DeleteMessageBatch", 69 | "sqs:SendMessage", 70 | "sqs:SendMessageBatch" 71 | ], 72 | "Resource": ["*"] 73 | } 74 | ] 75 | } 76 | ``` 77 | 78 | ## License 79 | 80 | dramatiq_sqs is licensed under Apache 2.0. Please see 81 | [LICENSE] for licensing details. 82 | 83 | 84 | [Dramatiq]: https://dramatiq.io 85 | [Amazon SQS]: https://aws.amazon.com/sqs/ 86 | [boto3]: https://boto3.readthedocs.io/en/latest/ 87 | [docs]: https://boto3.readthedocs.io/en/latest/guide/quickstart.html#configuration 88 | [LICENSE]: https://github.com/Bogdanp/dramatiq_sqs/blob/master/LICENSE 89 | [ElasticMQ]: https://github.com/adamw/elasticmq 90 | -------------------------------------------------------------------------------- /dramatiq_sqs/__init__.py: -------------------------------------------------------------------------------- 1 | from .broker import SQSBroker 2 | 3 | __version__ = "0.2.0" 4 | 5 | __all__ = [ 6 | "SQSBroker", 7 | "__version__", 8 | ] 9 | -------------------------------------------------------------------------------- /dramatiq_sqs/broker.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | from base64 import b64decode, b64encode 4 | from collections import deque 5 | from typing import Any, Dict, Iterable, List, Optional, Sequence, TypeVar 6 | 7 | import boto3 8 | import dramatiq 9 | from dramatiq.logging import get_logger 10 | 11 | #: The max number of bytes in a message. 12 | MAX_MESSAGE_SIZE = 256 * 1024 13 | 14 | #: The min and max number of seconds messages may be retained for. 15 | MIN_MESSAGE_RETENTION = 60 16 | MAX_MESSAGE_RETENTION = 14 * 86400 17 | 18 | #: The maximum amount of time SQS will wait before redelivering a 19 | #: message. This is also the maximum amount of time that a message can 20 | #: be delayed for. 21 | MAX_VISIBILITY_TIMEOUT = 2 * 3600 22 | 23 | #: The max number of messages that may be prefetched at a time. 24 | MAX_PREFETCH = 10 25 | 26 | #: The min value for WaitTimeSeconds. 27 | MIN_TIMEOUT = int(os.getenv("DRAMATIQ_SQS_MIN_TIMEOUT", "20")) 28 | 29 | #: The number of times a message will be received before being added 30 | #: to the dead-letter queue (if enabled). 31 | MAX_RECEIVES = 5 32 | 33 | 34 | class SQSBroker(dramatiq.Broker): 35 | """A Dramatiq_ broker that can be used with `Amazon SQS`_ 36 | 37 | This backend has a number of limitations compared to the built-in 38 | Redis and RMQ backends: 39 | 40 | * the max amount of time messages can be delayed by is 15 minutes, 41 | * messages can be at most 256KiB large, 42 | * messages must be processed within 2 hours of being pulled, 43 | otherwise they will be redelivered. 44 | 45 | The backend uses boto3_ under the hood. For details on how 46 | authorization works, check out its docs_. 47 | 48 | Parameters: 49 | namespace: The prefix to use when creating queues. 50 | middleware: The set of middleware that apply to this broker. 51 | retention: The number of seconds messages can be retained for. 52 | Defaults to 14 days. 53 | dead_letter: Whether to add a dead-letter queue. Defaults to false. 54 | max_receives: The number of times a message should be received before 55 | being added to the dead-letter queue. Defaults to MAX_RECEIVES. 56 | **options: Additional options that are passed to boto3. 57 | 58 | .. _Dramatiq: https://dramatiq.io 59 | .. _Amazon SQS: https://aws.amazon.com/sqs/ 60 | .. _boto3: http://boto3.readthedocs.io/en/latest/index.html 61 | .. _docs: http://boto3.readthedocs.io/en/latest/guide/configuration.html 62 | """ 63 | 64 | def __init__( 65 | self, *, 66 | namespace: Optional[str] = None, 67 | middleware: Optional[List[dramatiq.Middleware]] = None, 68 | retention: int = MAX_MESSAGE_RETENTION, 69 | dead_letter: bool = False, 70 | max_receives: int = MAX_RECEIVES, 71 | tags: Optional[Dict[str, str]] = None, 72 | **options, 73 | ) -> None: 74 | super().__init__(middleware=middleware) 75 | 76 | if retention < MIN_MESSAGE_RETENTION or retention > MAX_MESSAGE_RETENTION: 77 | raise ValueError(f"'retention' must be between {MIN_MESSAGE_RETENTION} and {MAX_MESSAGE_RETENTION}.") 78 | 79 | self.namespace: Optional[str] = namespace 80 | self.retention: str = str(retention) 81 | self.queues: Dict[str, Any] = {} 82 | self.dead_letter: bool = dead_letter 83 | self.max_receives: int = max_receives 84 | self.tags: Optional[Dict[str, str]] = tags 85 | self.sqs: Any = boto3.resource("sqs", **options) 86 | 87 | @property 88 | def consumer_class(self): 89 | return SQSConsumer 90 | 91 | def consume(self, queue_name: str, prefetch: int = 1, timeout: int = 30000) -> dramatiq.Consumer: 92 | try: 93 | return self.consumer_class(self.queues[queue_name], prefetch, timeout) 94 | except KeyError: # pragma: no cover 95 | raise dramatiq.QueueNotFound(queue_name) 96 | 97 | def declare_queue(self, queue_name: str) -> None: 98 | if queue_name not in self.queues: 99 | prefixed_queue_name = queue_name 100 | if self.namespace is not None: 101 | prefixed_queue_name = "%(namespace)s_%(queue_name)s" % { 102 | "namespace": self.namespace, 103 | "queue_name": queue_name, 104 | } 105 | 106 | self.emit_before("declare_queue", queue_name) 107 | 108 | self.queues[queue_name] = self._get_or_create_queue( 109 | QueueName=prefixed_queue_name, 110 | Attributes={ 111 | "MessageRetentionPeriod": self.retention, 112 | } 113 | ) 114 | if self.tags: 115 | self.sqs.meta.client.tag_queue( 116 | QueueUrl=self.queues[queue_name].url, 117 | Tags=self.tags 118 | ) 119 | 120 | if self.dead_letter: 121 | dead_letter_queue_name = f"{prefixed_queue_name}_dlq" 122 | dead_letter_queue = self._get_or_create_queue( 123 | QueueName=dead_letter_queue_name 124 | ) 125 | if self.tags: 126 | self.sqs.meta.client.tag_queue( 127 | QueueUrl=dead_letter_queue.url, 128 | Tags=self.tags 129 | ) 130 | redrive_policy = { 131 | "deadLetterTargetArn": dead_letter_queue.attributes["QueueArn"], 132 | "maxReceiveCount": str(self.max_receives) 133 | } 134 | self.queues[queue_name].set_attributes(Attributes={ 135 | "RedrivePolicy": json.dumps(redrive_policy) 136 | }) 137 | self.emit_after("declare_queue", queue_name) 138 | 139 | def _get_or_create_queue(self, **kwargs) -> Any: 140 | try: 141 | return self.sqs.get_queue_by_name(QueueName=kwargs['QueueName']) 142 | except self.sqs.meta.client.exceptions.QueueDoesNotExist: 143 | self.logger.debug(f'Queue does not exist, creating queue with params: {kwargs}') 144 | return self.sqs.create_queue(**kwargs) 145 | 146 | def enqueue(self, message: dramatiq.Message, *, delay: Optional[int] = None) -> dramatiq.Message: 147 | queue_name = message.queue_name 148 | if delay is None: 149 | queue = self.queues[queue_name] 150 | delay_seconds = 0 151 | elif delay <= 900000: 152 | queue = self.queues[queue_name] 153 | delay_seconds = int(delay / 1000) 154 | else: 155 | raise ValueError("Messages in SQS cannot be delayed for longer than 15 minutes.") 156 | 157 | encoded_message = b64encode(message.encode()).decode() 158 | if len(encoded_message) > MAX_MESSAGE_SIZE: 159 | raise RuntimeError("Messages in SQS can be at most 256KiB large.") 160 | 161 | self.logger.debug("Enqueueing message %r on queue %r.", message.message_id, queue_name) 162 | self.emit_before("enqueue", message, delay) 163 | queue.send_message( 164 | MessageBody=encoded_message, 165 | DelaySeconds=delay_seconds, 166 | ) 167 | self.emit_after("enqueue", message, delay) 168 | return message 169 | 170 | def get_declared_queues(self) -> Iterable[str]: 171 | return set(self.queues) 172 | 173 | def get_declared_delay_queues(self) -> Iterable[str]: 174 | return set() 175 | 176 | 177 | class SQSConsumer(dramatiq.Consumer): 178 | def __init__(self, queue: Any, prefetch: int, timeout: int) -> None: 179 | self.logger = get_logger(__name__, type(self)) 180 | self.queue = queue 181 | self.prefetch = min(prefetch, MAX_PREFETCH) 182 | self.visibility_timeout = MAX_VISIBILITY_TIMEOUT 183 | self.timeout = timeout # UNUSED 184 | self.messages: deque = deque() 185 | self.message_refc = 0 186 | 187 | def ack(self, message: "_SQSMessage") -> None: 188 | message._sqs_message.delete() 189 | self.message_refc -= 1 190 | 191 | #: Messages are added to DLQ by SQS redrive policy, so no actions are necessary 192 | nack = ack 193 | 194 | def requeue(self, messages: Iterable["_SQSMessage"]) -> None: 195 | for batch in chunk(messages, chunksize=10): 196 | # Re-enqueue batches of up to 10 messages. 197 | send_response = self.queue.send_messages(Entries=[{ 198 | "Id": str(i), 199 | "MessageBody": message._sqs_message.body, 200 | } for i, message in enumerate(batch)]) 201 | 202 | # Then delete the ones that were successfully re-enqueued. 203 | # The rest will have to wait until their visibility 204 | # timeout expires. 205 | failed_message_ids = [int(res["Id"]) for res in send_response.get("Failed", [])] 206 | requeued_messages = [m for i, m in enumerate(batch) if i not in failed_message_ids] 207 | self.queue.delete_messages(Entries=[{ 208 | "Id": str(i), 209 | "ReceiptHandle": message._sqs_message.receipt_handle, 210 | } for i, message in enumerate(requeued_messages)]) 211 | 212 | self.message_refc -= len(requeued_messages) 213 | 214 | def __next__(self) -> Optional[dramatiq.Message]: 215 | kw = { 216 | "MaxNumberOfMessages": self.prefetch, 217 | "WaitTimeSeconds": MIN_TIMEOUT, 218 | } 219 | if self.visibility_timeout is not None: 220 | kw["VisibilityTimeout"] = self.visibility_timeout 221 | 222 | try: 223 | return self.messages.popleft() 224 | except IndexError: 225 | if self.message_refc < self.prefetch: 226 | for sqs_message in self.queue.receive_messages(**kw): 227 | try: 228 | encoded_message = b64decode(sqs_message.body) 229 | dramatiq_message = dramatiq.Message.decode(encoded_message) 230 | self.messages.append(_SQSMessage(sqs_message, dramatiq_message)) 231 | self.message_refc += 1 232 | except Exception: # pragma: no cover 233 | self.logger.exception("Failed to decode message: %r", sqs_message.body) 234 | 235 | try: 236 | return self.messages.popleft() 237 | except IndexError: 238 | return None 239 | 240 | 241 | class _SQSMessage(dramatiq.MessageProxy): 242 | def __init__(self, sqs_message: Any, message: dramatiq.Message) -> None: 243 | super().__init__(message) 244 | 245 | self._sqs_message = sqs_message 246 | 247 | 248 | T = TypeVar("T") 249 | 250 | 251 | def chunk(xs: Iterable[T], *, chunksize=10) -> Iterable[Sequence[T]]: 252 | """Split a sequence into subseqs of chunksize length. 253 | """ 254 | chunk = [] 255 | for x in xs: 256 | chunk.append(x) 257 | if len(chunk) == chunksize: 258 | yield chunk 259 | chunk = [] 260 | 261 | if chunk: 262 | yield chunk 263 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [tool:pytest] 2 | testpaths = tests 3 | addopts = --cov dramatiq_sqs --cov-report html 4 | 5 | [pep8] 6 | max-line-length = 120 7 | 8 | [flake8] 9 | ignore = E402,F403,F811 10 | max-complexity = 20 11 | max-line-length = 120 12 | inline-quotes = double 13 | multiline-quotes = double 14 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from setuptools import setup 4 | 5 | 6 | def rel(*xs): 7 | return os.path.join(os.path.abspath(os.path.dirname(__file__)), *xs) 8 | 9 | 10 | with open(rel("dramatiq_sqs", "__init__.py"), "r") as f: 11 | version_marker = "__version__ = " 12 | for line in f: 13 | if line.startswith(version_marker): 14 | _, version = line.split(version_marker) 15 | version = version.strip().strip('"') 16 | break 17 | else: 18 | raise RuntimeError("Version marker not found.") 19 | 20 | setup( 21 | name="dramatiq_sqs", 22 | version=version, 23 | description="An Amazon SQS broker for Dramatiq.", 24 | long_description="Visit https://github.com/Bogdanp/dramatiq_sqs for more information.", 25 | packages=["dramatiq_sqs"], 26 | include_package_data=True, 27 | install_requires=["boto3", "dramatiq"], 28 | extras_require={ 29 | "dev": [ 30 | "bumpversion", 31 | "flake8", 32 | "flake8-quotes", 33 | "isort", 34 | "mypy", 35 | "pytest", 36 | "pytest-cov", 37 | "twine", 38 | ], 39 | }, 40 | python_requires=">=3.5", 41 | ) 42 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq_sqs/dad71a35dd71bfad737492473cc66399bfd28f73/tests/__init__.py -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import random 3 | import uuid 4 | 5 | import dramatiq 6 | import pytest 7 | from dramatiq.middleware import (AgeLimit, Callbacks, Pipelines, Retries, 8 | TimeLimit) 9 | 10 | from dramatiq_sqs import SQSBroker 11 | 12 | logfmt = "[%(asctime)s] [%(threadName)s] [%(name)s] [%(levelname)s] %(message)s" 13 | logging.basicConfig(level=logging.DEBUG, format=logfmt) 14 | logging.getLogger("botocore").setLevel(logging.WARN) 15 | random.seed(1337) 16 | 17 | 18 | @pytest.fixture 19 | def broker(): 20 | broker = SQSBroker( 21 | namespace="dramatiq_sqs_tests", 22 | middleware=[ 23 | AgeLimit(), 24 | TimeLimit(), 25 | Callbacks(), 26 | Pipelines(), 27 | Retries(min_backoff=1000, max_backoff=900000, max_retries=96), 28 | ], 29 | tags={ 30 | "owner": "dramatiq_sqs_tests", 31 | }, 32 | ) 33 | dramatiq.set_broker(broker) 34 | yield broker 35 | for queue in broker.queues.values(): 36 | queue.delete() 37 | 38 | 39 | @pytest.fixture 40 | def queue_name(broker): 41 | return f"queue_{uuid.uuid4()}" 42 | 43 | 44 | @pytest.fixture 45 | def worker(broker): 46 | worker = dramatiq.Worker(broker) 47 | worker.start() 48 | yield worker 49 | worker.stop() 50 | -------------------------------------------------------------------------------- /tests/test_broker.py: -------------------------------------------------------------------------------- 1 | import json 2 | import time 3 | 4 | import dramatiq 5 | import pytest 6 | from botocore.stub import Stubber 7 | 8 | from dramatiq_sqs import SQSBroker 9 | 10 | 11 | def test_can_enqueue_and_process_messages(broker, worker, queue_name): 12 | # Given that I have an actor that stores incoming messages in a database 13 | db = [] 14 | 15 | @dramatiq.actor(queue_name=queue_name) 16 | def do_work(x): 17 | db.append(x) 18 | 19 | # When I send that actor a message 20 | do_work.send(1) 21 | 22 | # And wait for it to be processed 23 | time.sleep(1) 24 | 25 | # Then the db should contain that message 26 | assert db == [1] 27 | 28 | 29 | def test_limits_prefetch_while_if_queue_is_full(broker, worker, queue_name): 30 | # Given that I have an actor that stores incoming messages in a database 31 | db = [] 32 | 33 | # Set the worker prefetch limit to 1 34 | worker.queue_prefetch = 1 35 | 36 | # Add delay to actor logic to simulate processing time 37 | @dramatiq.actor(queue_name=queue_name) 38 | def do_work(x): 39 | db.append(x) 40 | time.sleep(10) 41 | 42 | # When I send that actor messages, it'll only prefetch and process a single message 43 | do_work.send(1) 44 | do_work.send(2) 45 | 46 | # Wait for message to be processed 47 | time.sleep(2) 48 | 49 | # Then the db should contain only that message, while it sleeps 50 | assert db == [1] 51 | 52 | 53 | def test_can_enqueue_delayed_messages(broker, worker, queue_name): 54 | # Given that I have an actor that stores incoming messages in a database 55 | db = [] 56 | 57 | @dramatiq.actor(queue_name=queue_name) 58 | def do_work(x): 59 | db.append(x) 60 | 61 | # When I send that actor a delayed message 62 | start_time = time.time() 63 | do_work.send_with_options(args=(1,), delay=5000) 64 | 65 | # And poll the database for a result each second 66 | for _ in range(60): 67 | if db: 68 | break 69 | 70 | time.sleep(1) 71 | 72 | # Then the db should contain that message 73 | assert db == [1] 74 | 75 | # And an appropriate amount of time should have passed 76 | delta = time.time() - start_time 77 | assert delta >= 5 78 | 79 | 80 | def test_cant_delay_messages_for_longer_than_15_seconds(broker, queue_name): 81 | # Given that I have an actor 82 | @dramatiq.actor(queue_name=queue_name) 83 | def do_work(): 84 | pass 85 | 86 | # When I attempt to send that actor a message farther than 15 minutes into the future 87 | # Then I should get back a ValueError 88 | with pytest.raises(ValueError): 89 | do_work.send_with_options(delay=3600000) 90 | 91 | 92 | def test_cant_enqueue_messages_that_are_too_large(broker, queue_name): 93 | # Given that I have an actor 94 | @dramatiq.actor(queue_name=queue_name) 95 | def do_work(s): 96 | pass 97 | 98 | # When I attempt to send that actor a message that's too large 99 | # Then a RuntimeError should be raised 100 | with pytest.raises(RuntimeError): 101 | do_work.send("a" * 512 * 1024) 102 | 103 | 104 | def test_retention_period_is_validated(): 105 | # When I attempt to instantiate a broker with an invalid retention period 106 | # Then a ValueError should be raised 107 | with pytest.raises(ValueError): 108 | SQSBroker(retention=30 * 86400) 109 | 110 | 111 | def test_can_requeue_consumed_messages(broker, queue_name): 112 | # Given that I have an actor 113 | @dramatiq.actor(queue_name=queue_name) 114 | def do_work(): 115 | pass 116 | 117 | # When I send that actor a message 118 | do_work.send() 119 | 120 | # And consume the message off the queue 121 | consumer = broker.consume(queue_name) 122 | first_message = next(consumer) 123 | 124 | # And requeue the message 125 | consumer.requeue([first_message]) 126 | 127 | # Then I should be able to consume the message again immediately 128 | second_message = next(consumer) 129 | assert first_message == second_message 130 | 131 | 132 | def test_creates_dead_letter_queue(): 133 | # Given that I have an SQS broker with dead letters turned on 134 | broker = SQSBroker( 135 | namespace="dramatiq_sqs_tests", 136 | dead_letter=True, 137 | max_receives=20, 138 | ) 139 | 140 | # And I've stubbed out all the relevant API calls 141 | stubber = Stubber(broker.sqs.meta.client) 142 | error_response = { 143 | 'Error': { 144 | 'Code': 'AWS.SimpleQueueService.QueueDoesNotExist', 145 | 'Message': 'The specified queue does not exist.' 146 | } 147 | } 148 | stubber.add_client_error( 149 | "get_queue_url", 150 | http_status_code=404, 151 | service_error_code='QueueDoesNotExist', 152 | service_message='The specified queue does not exist.', 153 | response_meta=error_response) 154 | stubber.add_response("create_queue", {"QueueUrl": ""}) 155 | stubber.add_client_error( 156 | "get_queue_url", 157 | http_status_code=404, 158 | service_error_code='QueueDoesNotExist', 159 | service_message='The specified queue does not exist.', 160 | response_meta=error_response) 161 | 162 | stubber.add_response("create_queue", {"QueueUrl": ""}) 163 | stubber.add_response("get_queue_attributes", {"Attributes": {"QueueArn": "dlq"}}) 164 | stubber.add_response("set_queue_attributes", {}, { 165 | "QueueUrl": "", 166 | "Attributes": { 167 | "RedrivePolicy": json.dumps({ 168 | "deadLetterTargetArn": "dlq", 169 | "maxReceiveCount": "20" 170 | }) 171 | } 172 | }) 173 | 174 | # When I create a queue 175 | # Then a dead-letter queue should be created 176 | # And a redrive policy matching the queue and max receives should be added 177 | with stubber: 178 | broker.declare_queue("test") 179 | stubber.assert_no_pending_responses() 180 | 181 | 182 | def test_tags_queues_on_create(): 183 | # Given that I have an SQS broker with tags 184 | broker = SQSBroker( 185 | namespace="dramatiq_sqs_tests", 186 | tags={"key1": "value1", "key2": "value2"} 187 | ) 188 | 189 | # And I've stubbed out all the relevant API calls 190 | stubber = Stubber(broker.sqs.meta.client) 191 | error_response = { 192 | 'Error': { 193 | 'Code': 'AWS.SimpleQueueService.QueueDoesNotExist', 194 | 'Message': 'The specified queue does not exist.' 195 | } 196 | } 197 | stubber.add_client_error( 198 | "get_queue_url", 199 | http_status_code=404, 200 | service_error_code='QueueDoesNotExist', 201 | service_message='The specified queue does not exist.', 202 | response_meta=error_response) 203 | stubber.add_response("create_queue", {"QueueUrl": ""}) 204 | stubber.add_response("tag_queue", {}, { 205 | "QueueUrl": "", 206 | "Tags": { 207 | "key1": "value1", 208 | "key2": "value2" 209 | } 210 | }) 211 | 212 | # When I create a queue 213 | # Then the queue should have the specified tags 214 | with stubber: 215 | broker.declare_queue("test") 216 | stubber.assert_no_pending_responses() 217 | -------------------------------------------------------------------------------- /tests/test_utils.py: -------------------------------------------------------------------------------- 1 | from dramatiq_sqs.broker import chunk 2 | 3 | 4 | def test_chunk_can_split_iterators_into_chunks(): 5 | # Given that I have a range from 0 to 12 6 | xs = range(13) 7 | 8 | # When I pass that range to chunk with a chunksize of 2 9 | chunks = chunk(xs, chunksize=2) 10 | 11 | # Then I should get back these chunks 12 | assert list(chunks) == [ 13 | [0, 1], 14 | [2, 3], 15 | [4, 5], 16 | [6, 7], 17 | [8, 9], 18 | [10, 11], 19 | [12] 20 | ] 21 | --------------------------------------------------------------------------------