diff --git a/tests/service/test_amq_message_handler.py b/tests/service/test_amq_message_handler.py new file mode 100644 index 0000000..fb3097e --- /dev/null +++ b/tests/service/test_amq_message_handler.py @@ -0,0 +1,395 @@ +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.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()