From 34d87727fcba1ab54a5f6c449de6d05a5bdcc8d2 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Tue, 7 Apr 2026 12:33:52 +0100
Subject: [PATCH] Subscriber resilience: recreate consumer after connection
 failure

Move consumer creation from Subscriber.start() into the run() loop,
matching the pattern used by Consumer. If the connection drops and the
consumer is closed in the finally block, the loop now recreates it on
the next iteration instead of spinning forever on a None consumer.

Previously, start() created the consumer once and run() assumed it
existed for the lifetime of the subscriber. A connection failure would
set self.consumer = None in the finally cleanup, and the outer retry
loop would re-enter with no consumer, causing an infinite
'NoneType has no attribute receive' error loop.
---
 .../test_subscriber_graceful_shutdown.py      | 26 +++++++------------
 .../test_consumer_concurrency.py              |  3 +--
 trustgraph-base/trustgraph/base/subscriber.py | 23 ++++++++--------
 3 files changed, 22 insertions(+), 30 deletions(-)

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):