import asyncio import unittest from unittest.mock import AsyncMock, MagicMock, patch from amqp.adapter.backpressure_handler import BackpressureHandler from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter from amqp.adapter.data_message_factory import DataMessageFactory 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("1234567890abcdef1234") 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(MultipartDataMessage, result.__class__) 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 # [no ACK at this level, ACK is in parent call] self.mock_message.ack.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() # We verify that Message is created with correct content type by inspecting the RabbitMQ call args if available. self.assertEqual( "application/json", self.reply_to_exchange.publish.call_args[1]["message"].content_type ) if __name__ == "__main__": unittest.main()