735 lines
29 KiB
Python
735 lines
29 KiB
Python
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()
|