diff --git a/pyinfra/queue/queue_manager.py b/pyinfra/queue/queue_manager.py index 9644fb1..f226e0b 100644 --- a/pyinfra/queue/queue_manager.py +++ b/pyinfra/queue/queue_manager.py @@ -53,18 +53,23 @@ class QueueManager(object): callback = self._create_queue_callback(process_message_callback) self.logger.info("Consuming from queue") - consumer_tag = None + self._consumer_token = None try: - consumer_tag = self._channel.basic_consume(self._input_queue, callback) - self.logger.info(f"Registered with consumer-tag: {consumer_tag}") + self._consumer_token = self._channel.basic_consume(self._input_queue, callback) + self.logger.info(f"Registered with consumer-tag: {self._consumer_token}") self._channel.start_consuming() finally: self.logger.warning("An unhandled exception occurred while consuming messages. Consuming will stop.") - if consumer_tag: - self.logger.info(f"Cancelling subscription for consumer-tag: {consumer_tag}") - self._channel.basic_cancel(consumer_tag) - else: - self.logger.warning("No consumer-tag found, a possible subscription cannot be cancelled") + self.cancel_queue_subscription() + + def cancel_queue_subscription(self): + if self._consumer_token: + self.logger.info(f"Cancelling subscription for consumer-tag: {self._consumer_token}") + self._channel.basic_cancel(self._consumer_token) + + self._consumer_token = None + else: + self.logger.warning("No consumer-tag found, a possible subscription cannot be cancelled") def _create_queue_callback(self, process_message_callback: Callable): def callback(_channel, frame, properties, body):