amq-adapter-python/tools/test_amq_service.py

import unittest
from unittest.mock import AsyncMock, Mock, patch

from amqp.config.amq_configuration import AMQAdapter, AMQConfiguration
from amqp.service.amq_service import AMQService


class TestAMQService(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        # Create a mock configuration
        self.mock_adapter = Mock(spec=AMQAdapter)
        self.mock_adapter.service_name = "test-service"
        self.mock_adapter.route_mapping = {}
        self.mock_adapter.swarm_service_id = "test-swarm-id"
        self.mock_adapter.swarm_task_id = "test-task-id"
        self.mock_adapter.consul_host = "consul"
        self.mock_adapter.consul_port = 4500
        self.mock_adapter.consul_counter_key = "services/ids/counter"
        self.mock_adapter.consul_max_retries = 5
        self.mock_adapter.consul_initial_retry_delay = 0.2

        self.mock_config = Mock(spec=AMQConfiguration)
        self.mock_config.amq_adapter = self.mock_adapter

        # Add dispatch configuration
        self.mock_dispatch = Mock()
        self.mock_dispatch.amq_host = "localhost"
        self.mock_dispatch.amq_port = 5672
        self.mock_dispatch.rabbit_mq_user = "guest"
        self.mock_dispatch.rabbit_mq_password = "guest"
        self.mock_dispatch.use_dlq = False
        self.mock_config.dispatch = self.mock_dispatch

        # Add backpressure configuration
        self.mock_backpressure = Mock()
        self.mock_backpressure.threshold_load = 10
        self.mock_config.backpressure = self.mock_backpressure

        # Create service instance with mocked configuration
        self.service = AMQService(self.mock_config)

        # Mock the event loop
        self.mock_loop = Mock()
        self.service.loop = self.mock_loop

    def tearDown(self):
        # Clean up any resources
        if hasattr(self, "service"):
            self.service.cleanup()

    @patch("amqp.service.amq_service.RabbitMQClient")
    @patch("amqp.service.amq_service.DataMessageFactory")
    @patch("amqp.service.amq_service.ServiceMessageFactory")
    @patch("amqp.service.amq_service.initialize_telemetry")
    def test_init(
        self, mock_telemetry, mock_service_factory, mock_data_factory, mock_rabbit_client
    ):
        # Test initialization of AMQService
        service = AMQService(self.mock_config)

        # Verify that all required components are initialized
        self.assertIsNotNone(service.amq_configuration)
        self.assertIsNotNone(service.data_message_factory)
        self.assertIsNotNone(service.rabbit_mq_client)
        self.assertIsNotNone(service.tracer)
        self.assertIsNotNone(service.service_message_factory)

    @patch("amqp.service.amq_service.RabbitMQClient")
    async def test_register_routes_valid(self, mock_rabbit_client):
        # Setup mock route mapping
        self.mock_adapter.route_mapping = "test-service:test-exchange:test-queue:test.#:5:0"

        # Setup mock rabbit client
        mock_client = Mock()
        mock_client.register_inbound_routes = AsyncMock()
        mock_client.register_data_reply_callback = AsyncMock()
        mock_client.init = AsyncMock()
        mock_client.request_routes = AsyncMock()
        mock_client.channel = Mock()
        mock_client.channel.is_closed = False
        self.service.rabbit_mq_client = mock_client

        # Setup mock service adapter
        mock_service_adapter = Mock()
        mock_service_adapter.loop = self.mock_loop

        # Setup mock message handler
        mock_handler = Mock()
        mock_handler.inbound_data_message_callback_no_thread = AsyncMock()
        mock_handler.reply_received_callback = AsyncMock()
        self.service.amq_data_message_handler = mock_handler

        # Initialize the service
        await self.service.init(self.mock_loop, mock_service_adapter)

        # Call register_routes
        await self.service.register_routes()

        # Verify that routes were registered twice (once in init, once in register_routes)
        self.assertEqual(mock_client.register_inbound_routes.call_count, 2)
        self.assertEqual(mock_client.register_data_reply_callback.call_count, 2)

    @patch("amqp.service.amq_service.RabbitMQClient")
    async def test_register_routes_invalid(self, mock_rabbit_client):
        # Setup invalid route mapping
        self.mock_adapter.route_mapping = None

        # Setup mock rabbit client
        mock_client = Mock()
        mock_client.register_inbound_routes = AsyncMock()
        mock_client.register_data_reply_callback = AsyncMock()
        self.service.rabbit_mq_client = mock_client

        # Call register_routes
        result = await self.service.register_routes()

        # Verify that no routes were registered
        mock_client.register_inbound_routes.assert_not_called()
        mock_client.register_data_reply_callback.assert_not_called()
        self.assertIsNone(result)

    @patch("amqp.service.amq_service.RabbitMQClient")
    async def test_reinitialize_if_disconnected(self, mock_rabbit_client):
        # Setup mock route mapping
        self.mock_adapter.route_mapping = "test-service:test-exchange:test-queue:test.#:5:0"

        # Setup mock rabbit client with closed channel
        mock_client = AsyncMock()
        mock_client.channel = Mock()
        mock_client.channel.is_closed = True
        mock_client.init = AsyncMock()
        mock_client.register_inbound_routes = AsyncMock()
        mock_client.register_data_reply_callback = AsyncMock()
        mock_client.request_routes = AsyncMock()

        # Make the constructor return our mock client
        mock_rabbit_client.return_value = mock_client
        self.service.rabbit_mq_client = mock_client

        # Setup mock message handler
        mock_handler = Mock()
        mock_handler.inbound_data_message_callback_no_thread = AsyncMock()
        mock_handler.reply_received_callback = AsyncMock()
        self.service.amq_data_message_handler = mock_handler

        # Call reinitialize_if_disconnected
        await self.service.reinitialize_if_disconnected()

        # Verify that client was reinitialized
        mock_client.init.assert_called_once()

    def test_remove_outstanding_future(self):
        # Setup mock message handler with outstanding message
        self.service.amq_data_message_handler = Mock()
        self.service.amq_data_message_handler.outstanding = {"test-id": Mock()}

        # Remove future
        self.service.remove_outstanding_future("test-id")

        # Verify message was removed from outstanding
        self.assertNotIn("test-id", self.service.amq_data_message_handler.outstanding)

    def test_cleanup(self):
        # Setup mock rabbit client
        mock_client = Mock()
        self.service.rabbit_mq_client = mock_client

        # Call cleanup
        self.service.cleanup()

        # Verify rabbit client cleanup was called
        mock_client.cleanup.assert_called_once()