From c67c893dd37f4a17e2d5ee72f4b3da4437e33060 Mon Sep 17 00:00:00 2001 From: Matthias Bisping Date: Wed, 23 Mar 2022 17:03:20 +0100 Subject: [PATCH] inactivity_timeout for channel.consume --- pyinfra/queue/consumer.py | 4 ++-- pyinfra/queue/queue_manager/pika_queue_manager.py | 4 ++-- pyinfra/queue/queue_manager/queue_manager.py | 2 +- pyinfra/test/queue/queue_manager_mock.py | 2 +- pyinfra/test/unit_tests/consumer_test.py | 2 +- 5 files changed, 7 insertions(+), 7 deletions(-) diff --git a/pyinfra/queue/consumer.py b/pyinfra/queue/consumer.py index 6dd7d60..a446b85 100644 --- a/pyinfra/queue/consumer.py +++ b/pyinfra/queue/consumer.py @@ -9,5 +9,5 @@ class Consumer: def consume_and_publish(self): self.queue_manager.consume_and_publish(self.callback) - def consume(self): - yield from self.queue_manager.consume() + def consume(self, **kwargs): + yield from self.queue_manager.consume(**kwargs) diff --git a/pyinfra/queue/queue_manager/pika_queue_manager.py b/pyinfra/queue/queue_manager/pika_queue_manager.py index ca7f9c6..5544a10 100644 --- a/pyinfra/queue/queue_manager/pika_queue_manager.py +++ b/pyinfra/queue/queue_manager/pika_queue_manager.py @@ -117,9 +117,9 @@ class PikaQueueManager(QueueManager): def pull_request(self): return self.channel.basic_get(self._input_queue) - def consume(self): + def consume(self, inactivity_timeout=None): logger.debug("Consuming") - return self.channel.consume(self._input_queue) + return self.channel.consume(self._input_queue, inactivity_timeout=inactivity_timeout) def consume_and_publish(self, visitor): diff --git a/pyinfra/queue/queue_manager/queue_manager.py b/pyinfra/queue/queue_manager/queue_manager.py index a8af9d3..44b6baf 100644 --- a/pyinfra/queue/queue_manager/queue_manager.py +++ b/pyinfra/queue/queue_manager/queue_manager.py @@ -27,7 +27,7 @@ class QueueManager(abc.ABC): pass @abc.abstractmethod - def consume(self): + def consume(self, **kwargs): pass @abc.abstractmethod diff --git a/pyinfra/test/queue/queue_manager_mock.py b/pyinfra/test/queue/queue_manager_mock.py index 0f0cbd6..1928f12 100644 --- a/pyinfra/test/queue/queue_manager_mock.py +++ b/pyinfra/test/queue/queue_manager_mock.py @@ -22,7 +22,7 @@ class QueueManagerMock(QueueManager): def pull_request(self): return self._input_queue.popleft() - def consume(self): + def consume(self, **kwargs): while self._input_queue: yield self.pull_request() diff --git a/pyinfra/test/unit_tests/consumer_test.py b/pyinfra/test/unit_tests/consumer_test.py index 0846bad..436ead8 100644 --- a/pyinfra/test/unit_tests/consumer_test.py +++ b/pyinfra/test/unit_tests/consumer_test.py @@ -85,7 +85,7 @@ class TestConsumer: storage.put_object(**get_object_descriptor(message), data=gzip.compress(data)) queue_manager.publish_request(message) - requests = consumer.consume() + requests = consumer.consume(inactivity_timeout=5) for _, r in zip(items, requests): queue_manager.publish_response(r, visitor)