diff --git a/tests/unit/test_base/test_async_processor.py b/tests/unit/test_base/test_async_processor.py index 8e7ad70f..464e459a 100644 --- a/tests/unit/test_base/test_async_processor.py +++ b/tests/unit/test_base/test_async_processor.py @@ -14,19 +14,20 @@ from trustgraph.base.async_processor import AsyncProcessor class TestAsyncProcessorSimple(IsolatedAsyncioTestCase): """Test AsyncProcessor base class functionality""" - @patch('trustgraph.base.async_processor.PulsarClient') + @patch('trustgraph.base.async_processor.get_pubsub') @patch('trustgraph.base.async_processor.Consumer') @patch('trustgraph.base.async_processor.ProcessorMetrics') @patch('trustgraph.base.async_processor.ConsumerMetrics') - async def test_async_processor_initialization_basic(self, mock_consumer_metrics, mock_processor_metrics, - mock_consumer, mock_pulsar_client): + async def test_async_processor_initialization_basic(self, mock_consumer_metrics, mock_processor_metrics, + mock_consumer, mock_get_pubsub): """Test basic AsyncProcessor initialization""" # Arrange - mock_pulsar_client.return_value = MagicMock() + mock_backend = MagicMock() + mock_get_pubsub.return_value = mock_backend mock_consumer.return_value = MagicMock() mock_processor_metrics.return_value = MagicMock() mock_consumer_metrics.return_value = MagicMock() - + config = { 'id': 'test-async-processor', 'taskgroup': AsyncMock() @@ -42,14 +43,14 @@ class TestAsyncProcessorSimple(IsolatedAsyncioTestCase): assert processor.running == True assert hasattr(processor, 'config_handlers') assert processor.config_handlers == [] - - # Verify PulsarClient was created - mock_pulsar_client.assert_called_once_with(**config) - + + # Verify get_pubsub was called to create backend + mock_get_pubsub.assert_called_once_with(**config) + # Verify metrics were initialized mock_processor_metrics.assert_called_once() mock_consumer_metrics.assert_called_once() - + # Verify Consumer was created for config subscription mock_consumer.assert_called_once() diff --git a/tests/unit/test_base/test_publisher_graceful_shutdown.py b/tests/unit/test_base/test_publisher_graceful_shutdown.py index e15cb1ec..3c5cb967 100644 --- a/tests/unit/test_base/test_publisher_graceful_shutdown.py +++ b/tests/unit/test_base/test_publisher_graceful_shutdown.py @@ -8,22 +8,22 @@ from trustgraph.base.publisher import Publisher @pytest.fixture -def mock_pulsar_client(): - """Mock Pulsar client for testing.""" - client = MagicMock() +def mock_pulsar_backend(): + """Mock Pulsar backend for testing.""" + backend = MagicMock() producer = AsyncMock() producer.send = MagicMock() producer.flush = MagicMock() producer.close = MagicMock() - client.create_producer.return_value = producer - return client + backend.create_producer.return_value = producer + return backend @pytest.fixture -def publisher(mock_pulsar_client): +def publisher(mock_pulsar_backend): """Create Publisher instance for testing.""" return Publisher( - client=mock_pulsar_client, + backend=mock_pulsar_backend, topic="test-topic", schema=dict, max_size=10, @@ -34,12 +34,12 @@ def publisher(mock_pulsar_client): @pytest.mark.asyncio async def test_publisher_queue_drain(): """Verify Publisher drains queue on shutdown.""" - mock_client = MagicMock() + mock_backend = MagicMock() mock_producer = MagicMock() - mock_client.create_producer.return_value = mock_producer + mock_backend.create_producer.return_value = mock_producer publisher = Publisher( - client=mock_client, + backend=mock_backend, topic="test-topic", schema=dict, max_size=10, @@ -85,12 +85,12 @@ async def test_publisher_queue_drain(): @pytest.mark.asyncio async def test_publisher_rejects_messages_during_drain(): """Verify Publisher rejects new messages during shutdown.""" - mock_client = MagicMock() + mock_backend = MagicMock() mock_producer = MagicMock() - mock_client.create_producer.return_value = mock_producer + mock_backend.create_producer.return_value = mock_producer publisher = Publisher( - client=mock_client, + backend=mock_backend, topic="test-topic", schema=dict, max_size=10, @@ -113,12 +113,12 @@ async def test_publisher_rejects_messages_during_drain(): @pytest.mark.asyncio async def test_publisher_drain_timeout(): """Verify Publisher respects drain timeout.""" - mock_client = MagicMock() + mock_backend = MagicMock() mock_producer = MagicMock() - mock_client.create_producer.return_value = mock_producer + mock_backend.create_producer.return_value = mock_producer publisher = Publisher( - client=mock_client, + backend=mock_backend, topic="test-topic", schema=dict, max_size=10, @@ -169,12 +169,12 @@ async def test_publisher_drain_timeout(): @pytest.mark.asyncio async def test_publisher_successful_drain(): """Verify Publisher drains successfully under normal conditions.""" - mock_client = MagicMock() + mock_backend = MagicMock() mock_producer = MagicMock() - mock_client.create_producer.return_value = mock_producer + mock_backend.create_producer.return_value = mock_producer publisher = Publisher( - client=mock_client, + backend=mock_backend, topic="test-topic", schema=dict, max_size=10, @@ -224,12 +224,12 @@ async def test_publisher_successful_drain(): @pytest.mark.asyncio async def test_publisher_state_transitions(): """Test Publisher state transitions during graceful shutdown.""" - mock_client = MagicMock() + mock_backend = MagicMock() mock_producer = MagicMock() - mock_client.create_producer.return_value = mock_producer + mock_backend.create_producer.return_value = mock_producer publisher = Publisher( - client=mock_client, + backend=mock_backend, topic="test-topic", schema=dict, max_size=10, @@ -276,9 +276,9 @@ async def test_publisher_state_transitions(): @pytest.mark.asyncio async def test_publisher_exception_handling(): """Test Publisher handles exceptions during drain gracefully.""" - mock_client = MagicMock() + mock_backend = MagicMock() mock_producer = MagicMock() - mock_client.create_producer.return_value = mock_producer + mock_backend.create_producer.return_value = mock_producer # Mock producer.send to raise exception on second call call_count = 0 @@ -291,7 +291,7 @@ async def test_publisher_exception_handling(): mock_producer.send.side_effect = failing_send publisher = Publisher( - client=mock_client, + backend=mock_backend, topic="test-topic", schema=dict, max_size=10,