import asyncio import unittest from unittest.mock import AsyncMock, MagicMock, patch from pika.exchange_type import ExchangeType 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 ) async def test_send_command_success(self): """Test successful sending of a command message.""" # Setup route = MagicMock() route.exchange = "test-exchange" route.component_name = "test-component" message = MagicMock(spec=DataMessage) message.id.return_value = SnowflakeId.from_hex("1234567890abcdef1234") message.content_type.return_value = "application/json" # Mock the channel and exchange mock_channel = AsyncMock() mock_exchange = AsyncMock() self.rabbit_mq_client.channel = mock_channel mock_channel.declare_exchange.return_value = mock_exchange # Mock DataMessageFactory.serialize with patch( "amqp.adapter.data_message_factory.DataMessageFactory.serialize" ) as mock_serialize: mock_serialize.return_value = b"serialized-message" # Execute result = await self.handler.send_command(route, message) # Verify self.assertTrue(result) mock_serialize.assert_called_once_with(message) mock_channel.declare_exchange.assert_called_once_with( name="test-exchange", type=ExchangeType.topic ) mock_exchange.publish.assert_called_once() # Verify the published message call_args = mock_exchange.publish.call_args self.assertEqual(call_args[1]["routing_key"], "test.component") published_message = call_args[1]["message"] self.assertEqual(published_message.body, b"serialized-message") self.assertEqual(published_message.correlation_id, "1234567890abcdef1234") self.assertEqual(published_message.content_type, "application/json") async def test_send_command_content_type_list(self): """Test sending a command with content type as a list.""" # Setup route = MagicMock() route.exchange = "test-exchange" route.component_name = "test-component" message = MagicMock(spec=DataMessage) message.id.return_value = SnowflakeId.from_hex("1234567890abcdef1234") message.content_type.return_value = ["application/json", "text/plain"] # Mock the channel and exchange mock_channel = AsyncMock() mock_exchange = AsyncMock() self.rabbit_mq_client.channel = mock_channel mock_channel.declare_exchange.return_value = mock_exchange # Mock DataMessageFactory.serialize with patch( "amqp.adapter.data_message_factory.DataMessageFactory.serialize" ) as mock_serialize: mock_serialize.return_value = b"serialized-message" # Execute result = await self.handler.send_command(route, message) # Verify self.assertTrue(result) mock_exchange.publish.assert_called_once() # Verify the published message uses the first content type from the list call_args = mock_exchange.publish.call_args published_message = call_args[1]["message"] self.assertEqual(published_message.content_type, "application/json") async def test_send_rpc_success(self): """Test successful sending of an RPC message.""" # Setup route = MagicMock() route.exchange = "test-exchange" route.component_name = "test-component" message = MagicMock(spec=DataMessage) message.id.return_value = SnowflakeId.from_hex("1234567890abcdef1234") message.content_type.return_value = "application/json" # Mock the router to return a reply queue name self.rabbit_mq_client.router = MagicMock() self.rabbit_mq_client.router.get_reply_to_queue_name.return_value = "reply-queue" # Mock the channel and exchange mock_channel = AsyncMock() mock_exchange = AsyncMock() self.rabbit_mq_client.channel = mock_channel mock_channel.declare_exchange.return_value = mock_exchange # Mock DataMessageFactory.serialize with patch( "amqp.adapter.data_message_factory.DataMessageFactory.serialize" ) as mock_serialize: mock_serialize.return_value = b"serialized-message" # Execute future = await self.handler.send_rpc(route, message) # Verify self.assertIsInstance(future, asyncio.Future) self.assertFalse(future.done()) # Verify the future is stored in outstanding self.assertIn("1234567890abcdef1234", self.handler.outstanding) self.assertEqual(self.handler.outstanding["1234567890abcdef1234"], future) # Verify the message was published mock_serialize.assert_called_once_with(message) mock_channel.declare_exchange.assert_called_once_with( name="test-exchange", type=ExchangeType.topic ) mock_exchange.publish.assert_called_once() # Verify the published message call_args = mock_exchange.publish.call_args self.assertEqual(call_args[1]["routing_key"], "test.component") published_message = call_args[1]["message"] self.assertEqual(published_message.body, b"serialized-message") self.assertEqual(published_message.correlation_id, "1234567890abcdef1234") self.assertEqual(published_message.reply_to, "reply-queue") self.assertEqual(published_message.content_type, "application/json") async def test_send_rpc_content_type_list(self): """Test sending an RPC message with content type as a list.""" # Setup route = MagicMock() route.exchange = "test-exchange" route.component_name = "test-component" message = MagicMock(spec=DataMessage) message.id.return_value = SnowflakeId.from_hex("1234567890abcdef1234") message.content_type.return_value = ["application/json", "text/plain"] # Mock the router to return a reply queue name self.rabbit_mq_client.router = MagicMock() self.rabbit_mq_client.router.get_reply_to_queue_name.return_value = "reply-queue" # Mock the channel and exchange mock_channel = AsyncMock() mock_exchange = AsyncMock() self.rabbit_mq_client.channel = mock_channel mock_channel.declare_exchange.return_value = mock_exchange # Mock DataMessageFactory.serialize with patch( "amqp.adapter.data_message_factory.DataMessageFactory.serialize" ) as mock_serialize: mock_serialize.return_value = b"serialized-message" # Execute future = await self.handler.send_rpc(route, message) # Verify self.assertIsInstance(future, asyncio.Future) # Verify the published message uses the first content type from the list call_args = mock_exchange.publish.call_args published_message = call_args[1]["message"] self.assertEqual(published_message.content_type, "application/json") async def test_run_rpc_request_success(self): """Test successful processing of an RPC request.""" # Setup self.mock_data_message.magic.return_value = DataMessageMagicByte.RPC_REQUEST.value # Mock the on_rpc method self.handler.on_rpc = AsyncMock(return_value=self.mock_data_response) # Execute await self.handler.run_rpc_request(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.handler.on_rpc.assert_called_once() self.reply_to_exchange.publish.assert_called_once() self.backpressure_handler.decrease_current_load.assert_called_once() async def test_run_rpc_request_multipart(self): """Test processing an RPC multipart message with file cleanup.""" # Setup self.mock_data_message.magic.return_value = DataMessageMagicByte.RPC_REQUEST.value self.handler.on_rpc = AsyncMock(return_value=self.mock_data_response) multipart_message = MultipartDataMessage( self.mock_data_message, {"file1": "/tmp/file1.txt"} ) # Execute await self.handler.run_rpc_request(self.mock_message, multipart_message, self.loop) # Verify self.mock_os_remove.assert_called_once_with("/tmp/file1.txt") async def test_run_rpc_request_service_exception(self): """Test handling of an exception from the RPC service.""" # Setup self.mock_data_message.magic.return_value = DataMessageMagicByte.RPC_REQUEST.value self.handler.on_rpc = AsyncMock(side_effect=Exception("RPC service error")) # Execute await self.handler.run_rpc_request(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_rpc_request_reply_exception(self): """Test handling of an exception when sending the RPC reply.""" # Setup self.mock_data_message.magic.return_value = DataMessageMagicByte.RPC_REQUEST.value self.handler.on_rpc = AsyncMock(return_value=self.mock_data_response) self.reply_to_exchange.publish.side_effect = Exception("Reply error") # Execute await self.handler.run_rpc_request(self.mock_message, self.mock_data_message, self.loop) # Verify self.backpressure_handler.decrease_current_load.assert_called_once() async def test_run_rpc_request_file_removal_exception(self): """Test handling of an exception when removing files in RPC request.""" # Setup self.mock_data_message.magic.return_value = DataMessageMagicByte.RPC_REQUEST.value self.handler.on_rpc = AsyncMock(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_rpc_request(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_on_rpc_success(self): """Test successful RPC invocation.""" # Setup a_message = MagicMock() a_message.method.return_value = "test_package" a_message.domain.return_value = "TestClass" a_message.path.return_value = "test_method" a_message.id.return_value = SnowflakeId.from_hex("1234567890abcdef1234") a_message._base64body.decode.return_value = '{"param1": "value1"}' # Mock parse_map_string with patch("amqp.service.amq_message_handler.parse_map_string") as mock_parse: mock_parse.return_value = {"param1": "value1"} # Mock invoke_command with patch("amqp.service.amq_message_handler.invoke_command") as mock_invoke: mock_response = MagicMock(spec=DataResponse) mock_invoke.return_value = mock_response # Execute result = await self.handler.on_rpc(a_message) # Verify self.assertEqual(result, mock_response) mock_parse.assert_called_once_with('{"param1": "value1"}') mock_invoke.assert_called_once_with( "test_package", "TestClass", "test_method", "1234567890abcdef1234", {"param1": "value1"}, ) async def test_on_rpc_global_function(self): """Test RPC invocation of a global function.""" # Setup a_message = MagicMock() a_message.method.return_value = "test_package" a_message.domain.return_value = "__global__" a_message.path.return_value = "test_function" a_message.id.return_value = SnowflakeId.from_hex("1234567890abcdef1234") a_message._base64body.decode.return_value = '{"param1": "value1"}' # Mock parse_map_string with patch("amqp.service.amq_message_handler.parse_map_string") as mock_parse: mock_parse.return_value = {"param1": "value1"} # Mock invoke_command with patch("amqp.service.amq_message_handler.invoke_command") as mock_invoke: mock_response = MagicMock(spec=DataResponse) mock_invoke.return_value = mock_response # Execute result = await self.handler.on_rpc(a_message) # Verify self.assertEqual(result, mock_response) mock_parse.assert_called_once_with('{"param1": "value1"}') mock_invoke.assert_called_once_with( "test_package", "__global__", "test_function", "1234567890abcdef1234", {"param1": "value1"}, ) if __name__ == "__main__": unittest.main()