amq-adapter-python/tests/model/test_model.py

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()