amq-adapter-python/tests/service/test_amq_message_handler.py

import asyncio
import os
import time
import unittest
from unittest.mock import AsyncMock, MagicMock, patch, call

from aio_pika import Message
from opentelemetry.trace import SpanContext, TraceFlags

from amqp.adapter.backpressure_handler import BackpressureHandler
from amqp.adapter.cleverthis_service_adapter import AMQMessage, CleverThisServiceAdapter
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.data_response_factory import DataResponseFactory
from amqp.adapter.file_handler import FileHandler
from amqp.adapter.service_message_factory import ServiceMessageFactory
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import DataMessage, DataMessageMagicByte, DataResponse, MultipartDataMessage
from amqp.model.snowflake_id import SnowflakeId
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient
from amqp.service.amq_message_handler import DataMessageHandler, collect_trace_info, get_default_trace_state, set_text


class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        # Mock dependencies
        self.service_adapter = MagicMock(spec=CleverThisServiceAdapter)
        self.service_adapter.on_message = AsyncMock()
        
        self.tracer = MagicMock()
        self.tracer.start_span = MagicMock()
        
        self.message_factory = MagicMock(spec=DataMessageFactory)
        self.service_message_factory = MagicMock(spec=ServiceMessageFactory)
        self.reply_to_exchange = MagicMock()
        self.reply_to_exchange.publish = AsyncMock()
        self.reply_to_exchange.name = "test-exchange"
        
        self.rabbit_mq_client = MagicMock(spec=RabbitMQClient)
        self.rabbit_mq_client.connection = MagicMock()
        
        self.amq_configuration = MagicMock(spec=AMQConfiguration)
        self.amq_configuration.dispatch = MagicMock()
        self.amq_configuration.dispatch.rabbit_mq_url = "amqp://guest:guest@localhost:5672/"
        self.amq_configuration.dispatch.download_dir = "/tmp/test_downloads"
        
        self.loop = asyncio.get_event_loop()
        
        self.backpressure_handler = MagicMock(spec=BackpressureHandler)
        self.backpressure_handler.increase_current_load = MagicMock()
        self.backpressure_handler.decrease_current_load = MagicMock()
        self.backpressure_handler.check_overload_condition = AsyncMock()
        self.backpressure_handler.current_availability = 5
        self.backpressure_handler.max_availability = 10
        
        self.file_handler = MagicMock(spec=FileHandler)
        self.file_handler.on_message = AsyncMock(return_value={})
        
        # Create the handler
        self.handler = DataMessageHandler(
            service_adapter=self.service_adapter,
            tracer=self.tracer,
            message_factory=self.message_factory,
            service_message_factory=self.service_message_factory,
            reply_to_exchange=self.reply_to_exchange,
            rabbit_mq_client=self.rabbit_mq_client,
            amq_configuration=self.amq_configuration,
            loop=self.loop,
            backpressure_handler=self.backpressure_handler,
            file_handler=self.file_handler
        )
        
        # Create a mock message
        self.mock_message = MagicMock()
        self.mock_message.delivery_tag = 123
        self.mock_message.consumer_tag = "consumer-tag"
        self.mock_message.reply_to = "reply-queue"
        self.mock_message.correlation_id = "correlation-id"
        self.mock_message.headers = {"clevermicro.addressing": 0}
        self.mock_message.properties = {}
        self.mock_message.body = b"test-body"
        self.mock_message.ack = AsyncMock()
        self.mock_message.nack = AsyncMock()
        
        # Create a mock data message
        self.mock_data_message = MagicMock(spec=DataMessage)
        self.mock_data_message.magic.return_value = DataMessageMagicByte.HTTP_REQUEST.value
        self.mock_data_message.id.return_value = SnowflakeId.from_hex("1234567890abcdef")
        self.mock_data_message.method.return_value = "GET"
        self.mock_data_message.path.return_value = "/test/path"
        self.mock_data_message.trace_info.return_value = {"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"}
        self.mock_data_message.body.return_value = b'{"test": "data"}'
        
        # Create a mock data response
        self.mock_data_response = MagicMock(spec=DataResponse)
        self.mock_data_response.id.return_value = "response-id"
        self.mock_data_response.content_type.return_value = "application/json"
        
        # Patch DataMessageFactory.from_bytes
        self.from_bytes_patch = patch('amqp.adapter.data_message_factory.DataMessageFactory.from_bytes', 
                                      return_value=self.mock_data_message)
        self.mock_from_bytes = self.from_bytes_patch.start()
        
        # Patch DataMessageFactory.from_stream
        self.from_stream_patch = patch('amqp.adapter.data_message_factory.DataMessageFactory.from_stream', 
                                       return_value=self.mock_data_message)
        self.mock_from_stream = self.from_stream_patch.start()
        
        # Patch DataResponseFactory.serialize
        self.serialize_patch = patch('amqp.adapter.data_response_factory.DataResponseFactory.serialize', 
                                     return_value=b'serialized-response')
        self.mock_serialize = self.serialize_patch.start()
        
        # Patch DataResponseFactory.from_bytes
        self.response_from_bytes_patch = patch('amqp.adapter.data_response_factory.DataResponseFactory.from_bytes', 
                                              return_value=self.mock_data_response)
        self.mock_response_from_bytes = self.response_from_bytes_patch.start()
        
        # Patch asyncio.run_coroutine_threadsafe
        self.run_coroutine_patch = patch('asyncio.run_coroutine_threadsafe')
        self.mock_run_coroutine = self.run_coroutine_patch.start()
        
        # Patch time.time
        self.time_patch = patch('time.time', return_value=1234567890.0)
        self.mock_time = self.time_patch.start()
        
        # Patch os.remove
        self.os_remove_patch = patch('os.remove')
        self.mock_os_remove = self.os_remove_patch.start()
        
    def tearDown(self):
        self.from_bytes_patch.stop()
        self.from_stream_patch.stop()
        self.serialize_patch.stop()
        self.response_from_bytes_patch.stop()
        self.run_coroutine_patch.stop()
        self.time_patch.stop()
        self.os_remove_patch.stop()
    
    async def test_inbound_data_message_callback_no_thread_success(self):
        """Test successful processing of an inbound message."""
        # Setup
        self.service_adapter.on_message.return_value = self.mock_data_response
        
        # Execute
        await self.handler.inbound_data_message_callback_no_thread(self.mock_message)
        
        # Verify
        self.mock_from_bytes.assert_called_once_with(self.mock_message.body)
        self.tracer.start_span.assert_called_once()
        self.backpressure_handler.increase_current_load.assert_called_once()
        self.service_adapter.on_message.assert_called_once()
        self.reply_to_exchange.publish.assert_called_once()
        self.backpressure_handler.decrease_current_load.assert_called_once()
        self.mock_message.ack.assert_called_once_with(multiple=False)
    
    async def test_inbound_data_message_callback_no_thread_unknown_magic(self):
        """Test handling of a message with unknown magic value."""
        # Setup
        self.mock_data_message.magic.return_value = "UNKNOWN"
        
        # Execute
        await self.handler.inbound_data_message_callback_no_thread(self.mock_message)
        
        # Verify
        self.mock_message.nack.assert_called_once_with(requeue=False)
        self.mock_message.ack.assert_called_once_with(multiple=False)
    
    async def test_inbound_data_message_callback_no_thread_no_message(self):
        """Test handling when no valid message is present."""
        # Setup
        self.mock_from_bytes.return_value = None
        
        # Execute
        await self.handler.inbound_data_message_callback_no_thread(self.mock_message)
        
        # Verify
        self.mock_message.ack.assert_called_once_with(multiple=False)
    
    async def test_inbound_data_message_callback_as_thread(self):
        """Test the thread-based callback method."""
        # Setup
        with patch('threading.Thread') as mock_thread:
            mock_thread_instance = MagicMock()
            mock_thread.return_value = mock_thread_instance
            
            # Execute
            await self.handler.inbound_data_message_callback_as_thread(self.mock_message)
            
            # Verify
            mock_thread.assert_called_once()
            mock_thread_instance.start.assert_called_once()
    
    async def test_reconstitute_data_message_standard(self):
        """Test reconstituting a standard data message."""
        # Execute
        result = await self.handler.reconstitute_data_message(self.mock_message)
        
        # Verify
        self.assertEqual(result, self.mock_data_message)
        self.mock_from_bytes.assert_called_once_with(self.mock_message.body)
    
    async def test_reconstitute_data_message_stream(self):
        """Test reconstituting a message from a stream."""
        # Setup
        self.mock_message.headers = {"clevermicro.addressing": 1}
        
        # Execute
        result = await self.handler.reconstitute_data_message(self.mock_message)
        
        # Verify
        self.assertEqual(result, self.mock_data_message)
        self.mock_from_stream.assert_called_once()
        self.file_handler.on_message.assert_called_once()
    
    async def test_run_http_request_magic_success(self):
        """Test successful processing of an HTTP request."""
        # Setup
        self.service_adapter.on_message.return_value = self.mock_data_response
        
        # Execute
        await self.handler.run_http_request_magic(self.mock_message, self.mock_data_message, self.loop)
        
        # Verify
        self.backpressure_handler.increase_current_load.assert_called_once()
        self.backpressure_handler.check_overload_condition.assert_called()
        self.service_adapter.on_message.assert_called_once()
        self.reply_to_exchange.publish.assert_called_once()
        self.backpressure_handler.decrease_current_load.assert_called_once()
    
    async def test_run_http_request_magic_multipart(self):
        """Test processing a multipart message with file cleanup."""
        # Setup
        self.service_adapter.on_message.return_value = self.mock_data_response
        multipart_message = MultipartDataMessage(self.mock_data_message, {"file1": "/tmp/file1.txt"})
        
        # Execute
        await self.handler.run_http_request_magic(self.mock_message, multipart_message, self.loop)
        
        # Verify
        self.mock_os_remove.assert_called_once_with("/tmp/file1.txt")
    
    async def test_run_http_request_magic_service_exception(self):
        """Test handling of an exception from the service adapter."""
        # Setup
        self.service_adapter.on_message.side_effect = Exception("Service error")
        
        # Execute
        await self.handler.run_http_request_magic(self.mock_message, self.mock_data_message, self.loop)
        
        # Verify
        self.backpressure_handler.decrease_current_load.assert_called_once()
        self.backpressure_handler.check_overload_condition.assert_called()
    
    async def test_run_http_request_magic_reply_exception(self):
        """Test handling of an exception when sending the reply."""
        # Setup
        self.service_adapter.on_message.return_value = self.mock_data_response
        self.reply_to_exchange.publish.side_effect = Exception("Reply error")
        
        # Execute
        await self.handler.run_http_request_magic(self.mock_message, self.mock_data_message, self.loop)
        
        # Verify
        self.mock_message.nack.assert_called_once_with(requeue=False)
        self.backpressure_handler.decrease_current_load.assert_called_once()
    
    async def test_send_reply_success(self):
        """Test successful sending of a reply."""
        # Execute
        await self.handler.send_reply("reply-queue", self.mock_data_response)
        
        # Verify
        self.mock_run_coroutine.assert_called_once()
    
    async def test_send_reply_no_reply_to(self):
        """Test handling when no reply_to is provided."""
        # Execute
        await self.handler.send_reply(None, self.mock_data_response)
        
        # Verify
        self.mock_run_coroutine.assert_not_called()
    
    async def test_reply_received_callback_with_future(self):
        """Test handling a reply when there's a matching future."""
        # Setup
        future = asyncio.Future()
        self.handler.outstanding["correlation-id"] = future
        
        # Execute
        await self.handler.reply_received_callback(self.mock_message)
        
        # Verify
        self.assertTrue(future.done())
        self.assertEqual(future.result(), self.mock_data_response.body())
        self.mock_message.ack.assert_called_once_with(multiple=False)
    
    async def test_reply_received_callback_no_future(self):
        """Test handling a reply when there's no matching future."""
        # Execute
        await self.handler.reply_received_callback(self.mock_message)
        
        # Verify
        self.mock_message.ack.assert_called_once_with(multiple=False)
    
    def test_ensure_trace_info(self):
        """Test ensuring trace info is properly set."""
        # Execute
        self.handler.ensure_trace_info(self.mock_data_message)
        
        # Verify
        self.tracer.start_span.assert_called_once()
    
    def test_record_duration(self):
        """Test recording the duration of message processing."""
        # Setup
        start_time = 1234567880.0  # 10 seconds before mock_time
        
        # Execute
        self.handler.record_duration(start_time, 123)
        
        # No specific assertions needed as this is mostly logging
    
    def test_collect_trace_info(self):
        """Test collecting trace info."""
        # Execute
        result = collect_trace_info()
        
        # Verify it returns a dictionary
        self.assertIsInstance(result, dict)
    
    def test_get_default_trace_state(self):
        """Test getting default trace state."""
        # Execute
        result = get_default_trace_state()
        
        # Verify it returns a TraceState
        self.assertIsNotNone(result)
    
    def test_set_text(self):
        """Test setting text in a carrier."""
        # Setup
        carrier = {}
        
        # Execute
        set_text(carrier, "key", "value")
        
        # Verify
        self.assertEqual(carrier["key"], "value")
    
    
    async def test_reconstitute_data_message_stream_none_result(self):
        """Test reconstituting a message from a stream when the result is None."""
        # Setup
        self.mock_message.headers = {"clevermicro.addressing": 1}
        self.mock_from_stream.return_value = None
        
        # Execute
        result = await self.handler.reconstitute_data_message(self.mock_message)
        
        # Verify
        self.assertIsNone(result)
        self.mock_from_stream.assert_called_once()
        self.file_handler.on_message.assert_not_called()
    
    async def test_run_http_request_magic_file_removal_exception(self):
        """Test handling of an exception when removing files."""
        # Setup
        self.service_adapter.on_message.return_value = self.mock_data_response
        multipart_message = MultipartDataMessage(self.mock_data_message, {"file1": "/tmp/file1.txt"})
        self.mock_os_remove.side_effect = OSError("File removal error")
        
        # Execute
        await self.handler.run_http_request_magic(self.mock_message, multipart_message, self.loop)
        
        # Verify
        self.mock_os_remove.assert_called_once_with("/tmp/file1.txt")
        # Should continue without raising an exception
    
    async def test_send_reply_content_type_list(self):
        """Test sending a reply with content type as a list."""
        # Setup
        self.mock_data_response.content_type.return_value = ["application/json", "text/plain"]
        
        # Execute
        await self.handler.send_reply("reply-queue", self.mock_data_response)
        
        # Verify
        self.mock_run_coroutine.assert_called_once()
        # Check that the first content type was used
        call_args = self.mock_run_coroutine.call_args[0][0]
        # We verify that Message is created with correct content type by inspecting the call args if available.
        self.assertEqual(Message.call_args[1]['content_type'], "application/json")


if __name__ == "__main__":
    unittest.main()