import base64 import json import unittest from amqp.model.model import ( AMQErrorMessage, AMQRoute, CleverMicroMessage, DataMessage, DataResponse, ScalingRequest, ScalingRequestAlert, ) from amqp.model.snowflake_id import SnowflakeId class TestCleverMicroMessage(unittest.TestCase): def setUp(self): self.test_id = SnowflakeId() self.test_headers = {"Content-Type": ["application/json"]} self.test_trace_info = {"trace_id": "123"} self.test_body = b"test data" self.test_base64body = base64.b64encode(self.test_body) self.message = CleverMicroMessage( magic="A", id=self.test_id, m_field="GET", d_field="test_domain", p_field="test_path", headers=self.test_headers, trace_info=self.test_trace_info, base64body=self.test_base64body, ) def test_getters(self): self.assertEqual(self.message.magic(), "A") self.assertEqual(self.message.id(), self.test_id) self.assertEqual(self.message.m_field(), "GET") self.assertEqual(self.message.d_field(), "test_domain") self.assertEqual(self.message.p_field(), "test_path") self.assertEqual(self.message.headers(), self.test_headers) self.assertEqual(self.message.trace_info(), self.test_trace_info) self.assertEqual(self.message.base64body(), self.test_base64body) def test_body_decoding(self): self.assertEqual(self.message.body(), self.test_body) def test_content_type(self): self.assertEqual(self.message.content_type(), "application/json") # Test default content type message_no_content_type = CleverMicroMessage( magic="A", id=self.test_id, m_field="GET", d_field="test_domain", p_field="test_path", headers={}, trace_info={}, base64body=self.test_base64body, ) self.assertEqual( message_no_content_type.content_type(), CleverMicroMessage.DEFAULT_CONTENT_TYPE[0] ) class TestDataMessage(unittest.TestCase): def setUp(self): self.test_id = SnowflakeId() self.test_headers = {"Content-Type": ["application/json"]} self.test_trace_info = {"trace_id": "123"} self.test_base64body = base64.b64encode(b"test data") self.message = DataMessage( magic="A", id=self.test_id, method="GET", domain="testdomain", path="testpath", headers=self.test_headers, trace_info=self.test_trace_info, base64body=self.test_base64body, ) def test_getters(self): self.assertEqual(self.message.method(), "GET") self.assertEqual(self.message.domain(), "testdomain") self.assertEqual(self.message.path(), "testpath") def test_routing_key(self): self.assertEqual(self.message.routing_key(), "testdomain.testpath") class TestDataResponse(unittest.TestCase): def setUp(self): self.test_id = SnowflakeId() self.test_trace_info = {"trace_id": "123"} self.test_base64body = base64.b64encode(b"test data") self.response = DataResponse( magic="R", id=self.test_id, response_code="200", content_type="application/json", error="", error_cause="", headers={}, trace_info=self.test_trace_info, base64body=self.test_base64body, ) def test_getters(self): self.assertEqual(self.response.response_code(), "200") self.assertEqual(self.response.error(), "") self.assertEqual(self.response.error_cause(), "") self.assertEqual(self.response.content_type(), "application/json") class TestAMQRoute(unittest.TestCase): def setUp(self): self.route = AMQRoute( component_name="test_component", exchange="test_exchange", queue="test_queue", key="test_key", timeout=1000, valid_until=1735689600000, # Some future timestamp ) def test_str_representation(self): expected = "test_component:test_exchange:test_queue:test_key:1000:1735689600000" self.assertEqual(str(self.route), expected) def test_equality(self): same_route = AMQRoute( component_name="test_component", exchange="test_exchange", queue="test_queue", key="test_key", timeout=1000, valid_until=1735689600000, ) different_route = AMQRoute( component_name="other_component", exchange="test_exchange", queue="test_queue", key="test_key", timeout=1000, valid_until=1735689600000, ) self.assertEqual(self.route, same_route) self.assertNotEqual(self.route, different_route) class TestScalingRequest(unittest.TestCase): def setUp(self): self.request = ScalingRequest( service_id="test_service", task_id="test_task", max_availability=5, current_availability=3, request_type=ScalingRequestAlert.OVERLOAD, version=1, ) def test_to_json(self): json_str = self.request.to_json() data = json.loads(json_str) self.assertEqual(data["serviceId"], "test_service") self.assertEqual(data["taskId"], "test_task") self.assertEqual(data["maxAvailability"], 5) self.assertEqual(data["currentAvailability"], 3) self.assertEqual(data["requestType"], "OVERLOAD") self.assertEqual(data["version"], 1) def test_from_json(self): json_str = json.dumps( { "serviceId": "test_service", "taskId": "test_task", "maxAvailability": 5, "currentAvailability": 3, "requestType": "OVERLOAD", "version": 1, } ) request = ScalingRequest.from_json(json_str) self.assertEqual(request.service_id, "test_service") self.assertEqual(request.task_id, "test_task") self.assertEqual(request.max_availability, 5) self.assertEqual(request.current_availability, 3) self.assertEqual(request.request_type, ScalingRequestAlert.OVERLOAD) self.assertEqual(request.version, 1) class TestAMQErrorMessage(unittest.TestCase): def setUp(self): self.error = AMQErrorMessage( response_code=400, error="Bad Request", detail="Invalid input", content_type="application/json", ) def test_serialize(self): serialized = self.error.serialize() data = json.loads(serialized) self.assertEqual(data["response_code"], 400) self.assertEqual(data["error"], "Bad Request") self.assertEqual(data["detail"], "Invalid input") if __name__ == "__main__": unittest.main()