diff --git a/tests/unit/test_base/test_subscriber_graceful_shutdown.py b/tests/unit/test_base/test_subscriber_graceful_shutdown.py index 0587e3d6..ec14f66b 100644 --- a/tests/unit/test_base/test_subscriber_graceful_shutdown.py +++ b/tests/unit/test_base/test_subscriber_graceful_shutdown.py @@ -61,23 +61,21 @@ async def test_subscriber_deferred_acknowledgment_success(): max_size=10, backpressure_strategy="block" ) - - # Start subscriber to initialize consumer - await subscriber.start() - + subscriber.consumer = mock_consumer + # Create queue for subscription queue = await subscriber.subscribe("test-queue") - + # Create mock message with matching queue name msg = create_mock_message("test-queue", {"data": "test"}) - + # Process message await subscriber._process_message(msg) - + # Should acknowledge successful delivery mock_consumer.acknowledge.assert_called_once_with(msg) mock_consumer.negative_acknowledge.assert_not_called() - + # Message should be in queue assert not queue.empty() received_msg = await queue.get() @@ -108,9 +106,7 @@ async def test_subscriber_dropped_message_still_acks(): max_size=1, # Very small queue backpressure_strategy="drop_new" ) - - # Start subscriber to initialize consumer - await subscriber.start() + subscriber.consumer = mock_consumer # Create queue and fill it queue = await subscriber.subscribe("test-queue") @@ -151,9 +147,7 @@ async def test_subscriber_orphaned_message_acks(): max_size=10, backpressure_strategy="block" ) - - # Start subscriber to initialize consumer - await subscriber.start() + subscriber.consumer = mock_consumer # Don't create any queues - message will be orphaned # This simulates a response arriving after the waiter has unsubscribed @@ -189,9 +183,7 @@ async def test_subscriber_backpressure_strategies(): max_size=2, backpressure_strategy="drop_oldest" ) - - # Start subscriber to initialize consumer - await subscriber.start() + subscriber.consumer = mock_consumer queue = await subscriber.subscribe("test-queue") diff --git a/tests/unit/test_concurrency/test_consumer_concurrency.py b/tests/unit/test_concurrency/test_consumer_concurrency.py index 03244b73..55749920 100644 --- a/tests/unit/test_concurrency/test_consumer_concurrency.py +++ b/tests/unit/test_concurrency/test_consumer_concurrency.py @@ -294,9 +294,8 @@ class TestPollTimeout: raise type('Timeout', (Exception,), {})("timeout") mock_pulsar_consumer.receive = capture_receive - consumer.consumer = mock_pulsar_consumer - await consumer.consume_from_queue() + await consumer.consume_from_queue(mock_pulsar_consumer) assert received_kwargs.get("timeout_millis") == 100 diff --git a/trustgraph-base/trustgraph/base/subscriber.py b/trustgraph-base/trustgraph/base/subscriber.py index 36948131..4ed672bc 100644 --- a/trustgraph-base/trustgraph/base/subscriber.py +++ b/trustgraph-base/trustgraph/base/subscriber.py @@ -45,15 +45,6 @@ class Subscriber: async def start(self): - # Create consumer via backend - self.consumer = await asyncio.to_thread( - self.backend.create_consumer, - topic=self.topic, - subscription=self.subscription, - schema=self.schema, - consumer_type='exclusive', - ) - self.task = asyncio.create_task(self.run()) async def stop(self): @@ -80,6 +71,16 @@ class Subscriber: try: + # Create consumer if needed (first run or after failure) + if self.consumer is None: + self.consumer = await asyncio.to_thread( + self.backend.create_consumer, + topic=self.topic, + subscription=self.subscription, + schema=self.schema, + consumer_type='exclusive', + ) + if self.metrics: self.metrics.state("running") @@ -179,8 +180,8 @@ class Subscriber: if not self.running and not self.draining: return - - # If handler drops out, sleep a retry + + # Sleep before retry await asyncio.sleep(1) async def subscribe(self, id):