amq-adapter-python/tests/router/test_router_consumer.py

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

from aio_pika.abc import (
    AbstractIncomingMessage,
    AbstractRobustChannel,
    AbstractRobustQueue,
)

from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQRoute, ServiceMessage
from amqp.model.service_message_type import ServiceMessageType
from amqp.router.router_base import (
    CM_C_REQ_QUEUE,
    CM_REQUEST_EXCHANGE,
)
from amqp.router.router_consumer import RouterConsumer


class TestRouterConsumer(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        # Create a mock for amq_adapter
        self.amq_adapter = MagicMock()
        self.amq_adapter.adapter_prefix.return_value = "test_"
        self.amq_adapter.service_name = "test_service"

        # Create the config mock with amq_adapter
        self.config = MagicMock(spec=AMQConfiguration)
        self.config.amq_adapter = self.amq_adapter
        self.consumer = RouterConsumer(self.config)

        # Mock channel
        self.channel = AsyncMock(spec=AbstractRobustChannel)
        self.consumer.channel = self.channel

        # Mock queue
        self.queue = AsyncMock(spec=AbstractRobustQueue)
        self.channel.declare_queue.return_value = self.queue

    def test_init(self):
        """Test initialization of RouterConsumer"""
        self.assertIsNotNone(self.consumer)
        self.assertEqual(self.consumer.amq_configuration, self.config)

    @patch("amqp.router.router_consumer.logging_info")
    async def test_init_routing_paths_consumer_success(self, mock_logging):
        """Test successful initialization of routing paths consumer"""
        self.consumer.is_open = MagicMock(return_value=True)

        await self.consumer.init_routing_paths_consumer(self.channel)

        # Verify queue declaration
        self.channel.declare_queue.assert_called_once_with(
            name="test_" + CM_C_REQ_QUEUE, durable=True, exclusive=False, auto_delete=False
        )

        # Verify queue binding
        self.queue.bind.assert_called_once_with(exchange=CM_REQUEST_EXCHANGE, routing_key="")

        # Verify consumer setup
        self.queue.consume.assert_called_once()

    @patch("amqp.router.router_consumer.logging_error")
    async def test_init_routing_paths_consumer_channel_closed(self, mock_logging):
        """Test initialization when channel is closed"""
        self.consumer.is_open = MagicMock(return_value=False)

        await self.consumer.init_routing_paths_consumer(self.channel)

        mock_logging.assert_called_once()
        self.channel.declare_queue.assert_not_called()

    @patch("amqp.router.router_consumer.logging_debug")
    @patch("amqp.router.router_consumer.logging_info")
    async def test_handle_routing_data_request_message_valid(
        self, mock_logging_info, mock_logging_debug
    ):
        """Test handling of valid routing data request message"""
        # Setup
        message = AsyncMock(spec=AbstractIncomingMessage)
        message.delivery_tag = 1
        message.consumer_tag = "test_consumer"
        message.properties = MagicMock()
        message.body = b"test_body"

        # Mock service message factory
        mock_service_message = ServiceMessage(
            id=None,
            payload_type=0,
            message_type=ServiceMessageType.ROUTING_DATA_REQUEST,
            recipient_name="test_service",
            reply_to="other_service",
            trace_info={},
            message="test_payload",
        )
        self.consumer.service_message_factory = MagicMock()
        self.consumer.service_message_factory.from_bytes = MagicMock(
            return_value=mock_service_message
        )
        self.consumer.service_message_factory.of = MagicMock()
        self.consumer.wrap_own_routes = MagicMock(return_value="test_routes")
        self.consumer.publish_service_message = AsyncMock(return_value=True)

        # Execute
        await self.consumer.handle_routing_data_request_message(message)

        # Verify
        message.ack.assert_called_once_with(False)
        self.consumer.publish_service_message.assert_called_once()

    @patch("amqp.router.router_consumer.logging_info")
    async def test_register_route_success(self, mock_logging):
        """Test successful route registration"""
        # Setup
        route = AMQRoute(
            component_name="test_service",
            exchange="test_exchange",
            queue="test_queue",
            key="test_key",
            timeout=1000,
            valid_until=0,
        )
        self.consumer.is_open = MagicMock(return_value=True)
        self.consumer.publish_service_message = AsyncMock(return_value=True)
        self.consumer.add_own_route = MagicMock()

        # Execute
        await self.consumer.register_route(route)

        # Verify
        self.consumer.publish_service_message.assert_called_once()
        self.consumer.add_own_route.assert_called_once_with(route)

    @patch("amqp.router.router_consumer.logging_warning")
    async def test_register_route_channel_closed(self, mock_logging):
        """Test route registration when channel is closed"""
        # Setup
        route = AMQRoute(
            component_name="test_service",
            exchange="test_exchange",
            queue="test_queue",
            key="test_key",
            timeout=1000,
            valid_until=0,
        )
        self.consumer.is_open = MagicMock(return_value=False)

        # Execute
        await self.consumer.register_route(route)

        # Verify
        mock_logging.assert_called_once()

    @patch("amqp.router.router_consumer.logging_info")
    @patch("amqp.router.router_consumer.logging_warning")
    async def test_remove_route_on_cleanup(self, mock_logging_warning, mock_logging_info):
        """Test route removal during cleanup"""
        # Setup
        route = AMQRoute(
            component_name="test_service",
            exchange="test_exchange",
            queue="test_queue",
            key="test_key",
            timeout=1000,
            valid_until=0,
        )
        self.consumer.publish_service_message = AsyncMock(return_value=True)
        self.consumer.remove_route = MagicMock(return_value=True)

        # Execute
        result = await self.consumer.remove_route_on_cleanup(route)

        # Verify
        self.consumer.publish_service_message.assert_called_once()
        self.consumer.remove_route.assert_called_once_with(route)
        self.assertTrue(result)

    @patch("amqp.router.router_consumer.logging_info")
    @patch("amqp.router.router_consumer.logging_warning")
    def test_cleanup(self, mock_logging_warning, mock_logging_info):
        """Test cleanup process"""
        # Setup
        routes = {
            AMQRoute(
                component_name="test_service",
                exchange="test_exchange1",
                queue="test_queue1",
                key="test_key1",
                timeout=1000,
                valid_until=0,
            ),
            AMQRoute(
                component_name="test_service",
                exchange="test_exchange2",
                queue="test_queue2",
                key="test_key2",
                timeout=1000,
                valid_until=0,
            ),
        }
        self.consumer.get_routes = MagicMock(return_value=routes)
        self.consumer.remove_route_on_cleanup = AsyncMock(return_value=True)
        self.consumer.wrap_routes = MagicMock(return_value="test_routes")

        # Execute
        self.consumer.cleanup()

        # Verify
        self.assertEqual(self.consumer.remove_route_on_cleanup.call_count, 2)
        mock_logging_info.assert_called()