Adding tests

tests/unit/test_base/test_publisher_graceful_shutdown.py

@@ -0,0 +1,330 @@
+"""Unit tests for Publisher graceful shutdown functionality."""
+
+import pytest
+import asyncio
+import time
+from unittest.mock import AsyncMock, MagicMock, patch
+from trustgraph.base.publisher import Publisher
+
+
+@pytest.fixture
+def mock_pulsar_client():
+    """Mock Pulsar client for testing."""
+    client = MagicMock()
+    producer = AsyncMock()
+    producer.send = MagicMock()
+    producer.flush = MagicMock()
+    producer.close = MagicMock()
+    client.create_producer.return_value = producer
+    return client
+
+
+@pytest.fixture
+def publisher(mock_pulsar_client):
+    """Create Publisher instance for testing."""
+    return Publisher(
+        client=mock_pulsar_client,
+        topic="test-topic",
+        schema=dict,
+        max_size=10,
+        drain_timeout=2.0
+    )
+
+
+@pytest.mark.asyncio
+async def test_publisher_queue_drain():
+    """Verify Publisher drains queue on shutdown."""
+    mock_client = MagicMock()
+    mock_producer = MagicMock()
+    mock_client.create_producer.return_value = mock_producer
+    
+    publisher = Publisher(
+        client=mock_client,
+        topic="test-topic", 
+        schema=dict,
+        max_size=10,
+        drain_timeout=1.0  # Shorter timeout for testing
+    )
+    
+    # Don't start the actual run loop - just test the drain logic
+    # Fill queue with messages directly
+    for i in range(5):
+        await publisher.q.put((f"id-{i}", {"data": i}))
+    
+    # Verify queue has messages
+    assert not publisher.q.empty()
+    
+    # Mock the producer creation in run() method by patching
+    with patch.object(publisher, 'run') as mock_run:
+        # Create a realistic run implementation that processes the queue
+        async def mock_run_impl():
+            # Simulate the actual run logic for drain
+            producer = mock_producer
+            while not publisher.q.empty():
+                try:
+                    id, item = await asyncio.wait_for(publisher.q.get(), timeout=0.1)
+                    producer.send(item, {"id": id})
+                except asyncio.TimeoutError:
+                    break
+            producer.flush()
+            producer.close()
+        
+        mock_run.side_effect = mock_run_impl
+        
+        # Start and stop publisher
+        await publisher.start()
+        await publisher.stop()
+    
+    # Verify all messages were sent
+    assert publisher.q.empty()
+    assert mock_producer.send.call_count == 5
+    mock_producer.flush.assert_called_once()
+    mock_producer.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_publisher_rejects_messages_during_drain():
+    """Verify Publisher rejects new messages during shutdown."""
+    mock_client = MagicMock()
+    mock_producer = MagicMock()
+    mock_client.create_producer.return_value = mock_producer
+    
+    publisher = Publisher(
+        client=mock_client,
+        topic="test-topic",
+        schema=dict,
+        max_size=10,
+        drain_timeout=1.0
+    )
+    
+    # Don't start the actual run loop
+    # Add one message directly
+    await publisher.q.put(("id-1", {"data": 1}))
+    
+    # Start shutdown process manually
+    publisher.running = False
+    publisher.draining = True
+    
+    # Try to send message during drain
+    with pytest.raises(RuntimeError, match="Publisher is shutting down"):
+        await publisher.send("id-2", {"data": 2})
+
+
+@pytest.mark.asyncio
+async def test_publisher_drain_timeout():
+    """Verify Publisher respects drain timeout."""
+    mock_client = MagicMock()
+    mock_producer = MagicMock()
+    mock_client.create_producer.return_value = mock_producer
+    
+    publisher = Publisher(
+        client=mock_client,
+        topic="test-topic",
+        schema=dict,
+        max_size=10,
+        drain_timeout=0.2  # Short timeout for testing
+    )
+    
+    # Fill queue with many messages directly
+    for i in range(10):
+        await publisher.q.put((f"id-{i}", {"data": i}))
+    
+    # Mock slow message processing
+    def slow_send(*args, **kwargs):
+        time.sleep(0.1)  # Simulate slow send
+    
+    mock_producer.send.side_effect = slow_send
+    
+    with patch.object(publisher, 'run') as mock_run:
+        # Create a run implementation that respects timeout
+        async def mock_run_with_timeout():
+            producer = mock_producer
+            end_time = time.time() + publisher.drain_timeout
+            
+            while not publisher.q.empty() and time.time() < end_time:
+                try:
+                    id, item = await asyncio.wait_for(publisher.q.get(), timeout=0.05)
+                    producer.send(item, {"id": id})
+                except asyncio.TimeoutError:
+                    break
+            
+            producer.flush()
+            producer.close()
+        
+        mock_run.side_effect = mock_run_with_timeout
+        
+        start_time = time.time()
+        await publisher.start()
+        await publisher.stop()
+        end_time = time.time()
+    
+    # Should timeout quickly
+    assert end_time - start_time < 1.0
+    
+    # Should have called flush and close even with timeout
+    mock_producer.flush.assert_called_once()
+    mock_producer.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_publisher_successful_drain():
+    """Verify Publisher drains successfully under normal conditions."""
+    mock_client = MagicMock()
+    mock_producer = MagicMock()
+    mock_client.create_producer.return_value = mock_producer
+    
+    publisher = Publisher(
+        client=mock_client,
+        topic="test-topic",
+        schema=dict,
+        max_size=10,
+        drain_timeout=2.0
+    )
+    
+    # Add messages directly to queue
+    messages = []
+    for i in range(3):
+        msg = {"data": i}
+        await publisher.q.put((f"id-{i}", msg))
+        messages.append(msg)
+    
+    with patch.object(publisher, 'run') as mock_run:
+        # Create a successful drain implementation
+        async def mock_successful_drain():
+            producer = mock_producer
+            processed = []
+            
+            while not publisher.q.empty():
+                id, item = await publisher.q.get()
+                producer.send(item, {"id": id})
+                processed.append((id, item))
+            
+            producer.flush()
+            producer.close()
+            return processed
+        
+        mock_run.side_effect = mock_successful_drain
+        
+        await publisher.start()
+        await publisher.stop()
+    
+    # All messages should be sent
+    assert publisher.q.empty()
+    assert mock_producer.send.call_count == 3
+    
+    # Verify correct messages were sent
+    sent_calls = mock_producer.send.call_args_list
+    for i, call in enumerate(sent_calls):
+        args, kwargs = call
+        assert args[0] == {"data": i}  # message content
+        # Note: kwargs format depends on how send was called in mock
+        # Just verify message was sent with correct content
+
+
+@pytest.mark.asyncio
+async def test_publisher_state_transitions():
+    """Test Publisher state transitions during graceful shutdown."""
+    mock_client = MagicMock()
+    mock_producer = MagicMock()
+    mock_client.create_producer.return_value = mock_producer
+    
+    publisher = Publisher(
+        client=mock_client,
+        topic="test-topic",
+        schema=dict,
+        max_size=10,
+        drain_timeout=1.0
+    )
+    
+    # Initial state
+    assert publisher.running is True
+    assert publisher.draining is False
+    
+    # Add message directly
+    await publisher.q.put(("id-1", {"data": 1}))
+    
+    with patch.object(publisher, 'run') as mock_run:
+        # Mock run that simulates state transitions
+        async def mock_run_with_states():
+            # Simulate drain process
+            publisher.running = False
+            publisher.draining = True
+            
+            # Process messages
+            while not publisher.q.empty():
+                id, item = await publisher.q.get()
+                mock_producer.send(item, {"id": id})
+            
+            # Complete drain
+            publisher.draining = False
+            mock_producer.flush()
+            mock_producer.close()
+        
+        mock_run.side_effect = mock_run_with_states
+        
+        await publisher.start()
+        await publisher.stop()
+    
+    # Should have completed all state transitions
+    assert publisher.running is False
+    assert publisher.draining is False
+    mock_producer.send.assert_called_once()
+    mock_producer.flush.assert_called_once()
+    mock_producer.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_publisher_exception_handling():
+    """Test Publisher handles exceptions during drain gracefully."""
+    mock_client = MagicMock()
+    mock_producer = MagicMock()
+    mock_client.create_producer.return_value = mock_producer
+    
+    # Mock producer.send to raise exception on second call
+    call_count = 0
+    def failing_send(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if call_count == 2:
+            raise Exception("Send failed")
+    
+    mock_producer.send.side_effect = failing_send
+    
+    publisher = Publisher(
+        client=mock_client,
+        topic="test-topic",
+        schema=dict,
+        max_size=10,
+        drain_timeout=1.0
+    )
+    
+    # Add messages directly
+    await publisher.q.put(("id-1", {"data": 1}))
+    await publisher.q.put(("id-2", {"data": 2}))
+    
+    with patch.object(publisher, 'run') as mock_run:
+        # Mock run that handles exceptions gracefully
+        async def mock_run_with_exceptions():
+            producer = mock_producer
+            
+            while not publisher.q.empty():
+                try:
+                    id, item = await publisher.q.get()
+                    producer.send(item, {"id": id})
+                except Exception as e:
+                    # Log exception but continue processing
+                    continue
+            
+            # Always call flush and close
+            producer.flush()
+            producer.close()
+        
+        mock_run.side_effect = mock_run_with_exceptions
+        
+        await publisher.start()
+        await publisher.stop()
+    
+    # Should have attempted to send both messages
+    assert mock_producer.send.call_count == 2
+    mock_producer.flush.assert_called_once()
+    mock_producer.close.assert_called_once()

tests/unit/test_base/test_subscriber_graceful_shutdown.py

@@ -0,0 +1,315 @@
+"""Unit tests for Subscriber graceful shutdown functionality."""
+
+import pytest
+import asyncio
+import uuid
+from unittest.mock import AsyncMock, MagicMock, patch
+from trustgraph.base.subscriber import Subscriber
+
+
+@pytest.fixture
+def mock_pulsar_client():
+    """Mock Pulsar client for testing."""
+    client = MagicMock()
+    consumer = MagicMock()
+    consumer.receive = MagicMock()
+    consumer.acknowledge = MagicMock()
+    consumer.negative_acknowledge = MagicMock()
+    consumer.pause_message_listener = MagicMock()
+    consumer.unsubscribe = MagicMock()
+    consumer.close = MagicMock()
+    client.subscribe.return_value = consumer
+    return client
+
+
+@pytest.fixture
+def subscriber(mock_pulsar_client):
+    """Create Subscriber instance for testing."""
+    return Subscriber(
+        client=mock_pulsar_client,
+        topic="test-topic",
+        subscription="test-subscription",
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=10,
+        drain_timeout=2.0,
+        backpressure_strategy="block"
+    )
+
+
+def create_mock_message(message_id="test-id", data=None):
+    """Create a mock Pulsar message."""
+    msg = MagicMock()
+    msg.properties.return_value = {"id": message_id}
+    msg.value.return_value = data or {"test": "data"}
+    return msg
+
+
+@pytest.mark.asyncio
+async def test_subscriber_deferred_acknowledgment_success():
+    """Verify Subscriber only acks on successful delivery."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic",
+        subscription="test-subscription", 
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=10,
+        backpressure_strategy="block"
+    )
+    
+    # Create queue for subscription
+    queue = await subscriber.subscribe("test-queue")
+    
+    # Create mock message
+    msg = create_mock_message("msg-1", {"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()
+    assert received_msg == {"data": "test"}
+
+
+@pytest.mark.asyncio
+async def test_subscriber_deferred_acknowledgment_failure():
+    """Verify Subscriber negative acks on delivery failure."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic",
+        subscription="test-subscription",
+        consumer_name="test-consumer", 
+        schema=dict,
+        max_size=1,  # Very small queue
+        backpressure_strategy="drop_new"
+    )
+    
+    # Create queue and fill it
+    queue = await subscriber.subscribe("test-queue")
+    await queue.put({"existing": "data"})
+    
+    # Create mock message - should be dropped
+    msg = create_mock_message("msg-1", {"data": "test"})
+    
+    # Process message (should fail due to full queue + drop_new strategy)
+    await subscriber._process_message(msg)
+    
+    # Should negative acknowledge failed delivery
+    mock_consumer.negative_acknowledge.assert_called_once_with(msg)
+    mock_consumer.acknowledge.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_subscriber_backpressure_strategies():
+    """Test different backpressure strategies."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    # Test drop_oldest strategy
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic", 
+        subscription="test-subscription",
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=2,
+        backpressure_strategy="drop_oldest"
+    )
+    
+    queue = await subscriber.subscribe("test-queue")
+    
+    # Fill queue
+    await queue.put({"data": "old1"})
+    await queue.put({"data": "old2"})
+    
+    # Add new message (should drop oldest)
+    msg = create_mock_message("msg-1", {"data": "new"})
+    await subscriber._process_message(msg)
+    
+    # Should acknowledge delivery
+    mock_consumer.acknowledge.assert_called_once_with(msg)
+    
+    # Queue should have new message (old one dropped)
+    messages = []
+    while not queue.empty():
+        messages.append(await queue.get())
+    
+    # Should contain old2 and new (old1 was dropped)
+    assert len(messages) == 2
+    assert {"data": "new"} in messages
+
+
+@pytest.mark.asyncio
+async def test_subscriber_graceful_shutdown():
+    """Test Subscriber graceful shutdown with queue draining."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic",
+        subscription="test-subscription",
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=10,
+        drain_timeout=1.0
+    )
+    
+    await subscriber.start()
+    
+    # Create subscription with messages
+    queue = await subscriber.subscribe("test-queue")
+    await queue.put({"data": "msg1"})
+    await queue.put({"data": "msg2"})
+    
+    # Initial state
+    assert subscriber.running is True
+    assert subscriber.draining is False
+    
+    # Start shutdown
+    stop_task = asyncio.create_task(subscriber.stop())
+    
+    # Allow brief processing
+    await asyncio.sleep(0.1)
+    
+    # Should be in drain state
+    assert subscriber.running is False
+    assert subscriber.draining is True
+    
+    # Should pause message listener
+    mock_consumer.pause_message_listener.assert_called_once()
+    
+    # Complete shutdown
+    await stop_task
+    
+    # Should have cleaned up
+    mock_consumer.unsubscribe.assert_called_once()
+    mock_consumer.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_subscriber_drain_timeout():
+    """Test Subscriber respects drain timeout."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic",
+        subscription="test-subscription", 
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=10,
+        drain_timeout=0.1  # Very short timeout
+    )
+    
+    await subscriber.start()
+    
+    # Create subscription with many messages
+    queue = await subscriber.subscribe("test-queue")
+    for i in range(20):
+        await queue.put({"data": f"msg{i}"})
+    
+    import time
+    start_time = time.time()
+    await subscriber.stop()
+    end_time = time.time()
+    
+    # Should timeout quickly
+    assert end_time - start_time < 1.0
+    
+    # Queue should still have messages (drain timed out)
+    assert not queue.empty()
+
+
+@pytest.mark.asyncio
+async def test_subscriber_pending_acks_cleanup():
+    """Test Subscriber cleans up pending acknowledgments on shutdown."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic",
+        subscription="test-subscription",
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=10
+    )
+    
+    await subscriber.start()
+    
+    # Add pending acknowledgments manually (simulating in-flight messages)
+    msg1 = create_mock_message("msg-1")
+    msg2 = create_mock_message("msg-2") 
+    subscriber.pending_acks["ack-1"] = msg1
+    subscriber.pending_acks["ack-2"] = msg2
+    
+    # Stop subscriber
+    await subscriber.stop()
+    
+    # Should negative acknowledge pending messages
+    assert mock_consumer.negative_acknowledge.call_count == 2
+    mock_consumer.negative_acknowledge.assert_any_call(msg1)
+    mock_consumer.negative_acknowledge.assert_any_call(msg2)
+    
+    # Pending acks should be cleared
+    assert len(subscriber.pending_acks) == 0
+
+
+@pytest.mark.asyncio
+async def test_subscriber_multiple_subscribers():
+    """Test Subscriber with multiple concurrent subscribers."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic",
+        subscription="test-subscription",
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=10
+    )
+    
+    # Create multiple subscriptions
+    queue1 = await subscriber.subscribe("queue-1")
+    queue2 = await subscriber.subscribe("queue-2")
+    queue_all = await subscriber.subscribe_all("queue-all")
+    
+    # Process message
+    msg = create_mock_message("msg-1", {"data": "broadcast"})
+    await subscriber._process_message(msg)
+    
+    # Should acknowledge (successful delivery to all queues)
+    mock_consumer.acknowledge.assert_called_once_with(msg)
+    
+    # Message should be in specific queue (queue-1) and broadcast queue
+    assert not queue1.empty()
+    assert queue2.empty()  # No message for queue-2
+    assert not queue_all.empty()
+    
+    # Verify message content
+    msg1 = await queue1.get()
+    msg_all = await queue_all.get()
+    assert msg1 == {"data": "broadcast"}
+    assert msg_all == {"data": "broadcast"}