feat: add polling_wait_seconds kwarg on consume_messages (#14)

allisson · Apr 12, 2024 · 78f1d9d · 78f1d9d
1 parent ed6e291
commit 78f1d9d
Show file tree

Hide file tree

Showing 6 changed files with 335 additions and 279 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v4.6.0
     hooks:
       - id: check-ast
       - id: fix-byte-order-marker
@@ -18,7 +18,7 @@ repos:
       - id: trailing-whitespace
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.2.2
+    rev: v0.3.7
     hooks:
       - id: ruff
       - id: ruff-format
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sqsx"
-version = "0.3.1"
+version = "0.5.0"
 description = "A simple task processor for Amazon SQS"
 authors = ["Allisson Azevedo"]
 readme = "README.md"

diff --git a/setup.py b/setup.py
@@ -7,7 +7,7 @@
 
 setup(
     name="sqsx",
-    version="0.4.0",
+    version="0.5.0",
     description="A simple task processor for Amazon SQS",
     long_description=long_description,
     long_description_content_type="text/markdown",

diff --git a/sqsx/queue.py b/sqsx/queue.py
@@ -15,7 +15,12 @@
 
 class BaseQueueMixin:
     def consume_messages(
-        self, max_messages: int = 1, max_threads: int = 1, wait_seconds: int = 10, run_forever: bool = True
+        self,
+        max_messages: int = 1,
+        max_threads: int = 1,
+        wait_seconds: int = 10,
+        polling_wait_seconds: int = 10,
+        run_forever: bool = True,
     ) -> None:
         logger.info(f"Starting consuming tasks, queue_url={self.url}")
         signal.signal(signal.SIGINT, self._exit_gracefully)
@@ -31,6 +36,7 @@ def consume_messages(
                 AttributeNames=["All"],
                 MaxNumberOfMessages=min(max_messages, 10),
                 MessageAttributeNames=["All"],
+                WaitTimeSeconds=polling_wait_seconds,
             )
 
             sqs_messages = response.get("Messages", [])

diff --git a/tests/test_queue.py b/tests/test_queue.py
@@ -216,7 +216,7 @@ def test_queue_exit_gracefully(queue):
     queue.add_task_handler("my_task", handler)
     queue.add_task("my_task", a=1, b=2, c=3)
 
-    queue.consume_messages(wait_seconds=1, run_forever=True)
+    queue.consume_messages(wait_seconds=1, polling_wait_seconds=0, run_forever=True)
 
     assert handler.result_sum == 6
 
@@ -291,6 +291,6 @@ def test_raw_queue_exit_gracefully(raw_queue):
     raw_queue.add_message(message_body="Message Body")
     raw_queue.add_message(message_body="Message Body")
 
-    raw_queue.consume_messages(wait_seconds=1, run_forever=True)
+    raw_queue.consume_messages(wait_seconds=1, polling_wait_seconds=0, run_forever=True)
 
     assert handler.call_count == 3