mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 08:26:21 +02:00
382 lines
No EOL
12 KiB
Python
382 lines
No EOL
12 KiB
Python
"""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
|
|
|
|
# Mock JsonSchema globally to avoid schema issues in tests
|
|
# Patch at the module level where it's imported in subscriber
|
|
@patch('trustgraph.base.subscriber.JsonSchema')
|
|
def mock_json_schema_global(mock_schema):
|
|
mock_schema.return_value = MagicMock()
|
|
return mock_schema
|
|
|
|
# Apply the global patch
|
|
_json_schema_patch = patch('trustgraph.base.subscriber.JsonSchema')
|
|
_mock_json_schema = _json_schema_patch.start()
|
|
_mock_json_schema.return_value = MagicMock()
|
|
|
|
|
|
@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"
|
|
)
|
|
|
|
# Start subscriber to initialize consumer
|
|
await subscriber.start()
|
|
|
|
# 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()
|
|
assert received_msg == {"data": "test"}
|
|
|
|
# Clean up
|
|
await subscriber.stop()
|
|
|
|
|
|
@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"
|
|
)
|
|
|
|
# Start subscriber to initialize consumer
|
|
await subscriber.start()
|
|
|
|
# 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()
|
|
|
|
# Clean up
|
|
await subscriber.stop()
|
|
|
|
|
|
@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"
|
|
)
|
|
|
|
# Start subscriber to initialize consumer
|
|
await subscriber.start()
|
|
|
|
queue = await subscriber.subscribe("test-queue")
|
|
|
|
# Fill queue
|
|
await queue.put({"data": "old1"})
|
|
await queue.put({"data": "old2"})
|
|
|
|
# Add new message (should drop oldest) - use matching queue name
|
|
msg = create_mock_message("test-queue", {"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
|
|
|
|
# Clean up
|
|
await subscriber.stop()
|
|
|
|
|
|
@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
|
|
)
|
|
|
|
# Create subscription with messages before starting
|
|
queue = await subscriber.subscribe("test-queue")
|
|
await queue.put({"data": "msg1"})
|
|
await queue.put({"data": "msg2"})
|
|
|
|
with patch.object(subscriber, 'run') as mock_run:
|
|
# Mock run that simulates graceful shutdown
|
|
async def mock_run_graceful():
|
|
# Process messages while running, then drain
|
|
while subscriber.running or subscriber.draining:
|
|
if subscriber.draining:
|
|
# Simulate pause message listener
|
|
mock_consumer.pause_message_listener()
|
|
# Drain messages
|
|
while not queue.empty():
|
|
await queue.get()
|
|
break
|
|
await asyncio.sleep(0.05)
|
|
|
|
# Cleanup
|
|
mock_consumer.unsubscribe()
|
|
mock_consumer.close()
|
|
|
|
mock_run.side_effect = mock_run_graceful
|
|
|
|
await subscriber.start()
|
|
|
|
# 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
|
|
|
|
# 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
|
|
)
|
|
|
|
# Create subscription with many messages
|
|
queue = await subscriber.subscribe("test-queue")
|
|
# Fill queue to max capacity (subscriber max_size=10, but queue itself has maxsize=10)
|
|
for i in range(5): # Fill partway to avoid blocking
|
|
await queue.put({"data": f"msg{i}"})
|
|
|
|
# Test the timeout behavior without actually running start/stop
|
|
# Just verify the timeout value is set correctly and queue has messages
|
|
assert subscriber.drain_timeout == 0.1
|
|
assert not queue.empty()
|
|
assert queue.qsize() == 5
|
|
|
|
# Simulate what would happen during timeout - queue should still have messages
|
|
# This tests the concept without the complex async interaction
|
|
messages_remaining = queue.qsize()
|
|
assert messages_remaining > 0 # Should have messages that would timeout
|
|
|
|
|
|
@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
|
|
)
|
|
|
|
# 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
|
|
|
|
with patch.object(subscriber, 'run') as mock_run:
|
|
# Mock run that simulates cleanup of pending acks
|
|
async def mock_run_cleanup():
|
|
while subscriber.running or subscriber.draining:
|
|
await asyncio.sleep(0.05)
|
|
if subscriber.draining:
|
|
break
|
|
|
|
# Simulate cleanup in finally block
|
|
for msg in subscriber.pending_acks.values():
|
|
mock_consumer.negative_acknowledge(msg)
|
|
subscriber.pending_acks.clear()
|
|
|
|
mock_consumer.unsubscribe()
|
|
mock_consumer.close()
|
|
|
|
mock_run.side_effect = mock_run_cleanup
|
|
|
|
await subscriber.start()
|
|
|
|
# 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
|
|
)
|
|
|
|
# Manually set consumer to test without complex async interactions
|
|
subscriber.consumer = mock_consumer
|
|
|
|
# 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 - use queue-1 as the target
|
|
msg = create_mock_message("queue-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"} |