[kimi] refactor: extract WebSocket message types into shared protocol module (#667) (#696)

tests/unit/test_protocol.py

@@ -0,0 +1,173 @@
+"""Tests for infrastructure.protocol — WebSocket message types."""
+
+import json
+
+import pytest
+
+from infrastructure.protocol import (
+    AgentStateMessage,
+    BarkMessage,
+    ConnectionAckMessage,
+    ErrorMessage,
+    MemoryFlashMessage,
+    MessageType,
+    SystemStatusMessage,
+    TaskUpdateMessage,
+    ThoughtMessage,
+    VisitorStateMessage,
+    WSMessage,
+)
+
+# ---------------------------------------------------------------------------
+# MessageType enum
+# ---------------------------------------------------------------------------
+
+
+class TestMessageType:
+    """MessageType enum covers all 9 Matrix PROTOCOL.md types."""
+
+    def test_has_all_nine_types(self):
+        assert len(MessageType) == 9
+
+    @pytest.mark.parametrize(
+        "member,value",
+        [
+            (MessageType.AGENT_STATE, "agent_state"),
+            (MessageType.VISITOR_STATE, "visitor_state"),
+            (MessageType.BARK, "bark"),
+            (MessageType.THOUGHT, "thought"),
+            (MessageType.SYSTEM_STATUS, "system_status"),
+            (MessageType.CONNECTION_ACK, "connection_ack"),
+            (MessageType.ERROR, "error"),
+            (MessageType.TASK_UPDATE, "task_update"),
+            (MessageType.MEMORY_FLASH, "memory_flash"),
+        ],
+    )
+    def test_enum_values(self, member, value):
+        assert member.value == value
+
+    def test_str_comparison(self):
+        """MessageType is a str enum so it can be compared to plain strings."""
+        assert MessageType.BARK == "bark"
+
+
+# ---------------------------------------------------------------------------
+# to_json / from_json round-trip
+# ---------------------------------------------------------------------------
+
+
+class TestAgentStateMessage:
+    def test_defaults(self):
+        msg = AgentStateMessage()
+        assert msg.type == "agent_state"
+        assert msg.agent_id == ""
+        assert msg.data == {}
+
+    def test_round_trip(self):
+        msg = AgentStateMessage(agent_id="timmy", data={"mood": "happy"}, ts=1000.0)
+        raw = msg.to_json()
+        restored = AgentStateMessage.from_json(raw)
+        assert restored.agent_id == "timmy"
+        assert restored.data == {"mood": "happy"}
+        assert restored.ts == 1000.0
+
+    def test_to_json_structure(self):
+        msg = AgentStateMessage(agent_id="timmy", data={"x": 1}, ts=123.0)
+        parsed = json.loads(msg.to_json())
+        assert parsed["type"] == "agent_state"
+        assert parsed["agent_id"] == "timmy"
+        assert parsed["data"] == {"x": 1}
+        assert parsed["ts"] == 123.0
+
+
+class TestVisitorStateMessage:
+    def test_round_trip(self):
+        msg = VisitorStateMessage(visitor_id="v1", data={"page": "/"}, ts=1.0)
+        restored = VisitorStateMessage.from_json(msg.to_json())
+        assert restored.visitor_id == "v1"
+        assert restored.data == {"page": "/"}
+
+
+class TestBarkMessage:
+    def test_round_trip(self):
+        msg = BarkMessage(agent_id="timmy", content="woof!", ts=1.0)
+        restored = BarkMessage.from_json(msg.to_json())
+        assert restored.agent_id == "timmy"
+        assert restored.content == "woof!"
+
+
+class TestThoughtMessage:
+    def test_round_trip(self):
+        msg = ThoughtMessage(agent_id="timmy", content="hmm...", ts=1.0)
+        restored = ThoughtMessage.from_json(msg.to_json())
+        assert restored.content == "hmm..."
+
+
+class TestSystemStatusMessage:
+    def test_round_trip(self):
+        msg = SystemStatusMessage(status="healthy", data={"uptime": 3600}, ts=1.0)
+        restored = SystemStatusMessage.from_json(msg.to_json())
+        assert restored.status == "healthy"
+        assert restored.data == {"uptime": 3600}
+
+
+class TestConnectionAckMessage:
+    def test_round_trip(self):
+        msg = ConnectionAckMessage(client_id="abc-123", ts=1.0)
+        restored = ConnectionAckMessage.from_json(msg.to_json())
+        assert restored.client_id == "abc-123"
+
+
+class TestErrorMessage:
+    def test_round_trip(self):
+        msg = ErrorMessage(code="INVALID", message="bad request", ts=1.0)
+        restored = ErrorMessage.from_json(msg.to_json())
+        assert restored.code == "INVALID"
+        assert restored.message == "bad request"
+
+
+class TestTaskUpdateMessage:
+    def test_round_trip(self):
+        msg = TaskUpdateMessage(task_id="t1", status="completed", data={"result": "ok"}, ts=1.0)
+        restored = TaskUpdateMessage.from_json(msg.to_json())
+        assert restored.task_id == "t1"
+        assert restored.status == "completed"
+        assert restored.data == {"result": "ok"}
+
+
+class TestMemoryFlashMessage:
+    def test_round_trip(self):
+        msg = MemoryFlashMessage(agent_id="timmy", memory_key="fav_food", content="kibble", ts=1.0)
+        restored = MemoryFlashMessage.from_json(msg.to_json())
+        assert restored.memory_key == "fav_food"
+        assert restored.content == "kibble"
+
+
+# ---------------------------------------------------------------------------
+# WSMessage.from_json dispatch
+# ---------------------------------------------------------------------------
+
+
+class TestWSMessageDispatch:
+    """WSMessage.from_json dispatches to the correct subclass."""
+
+    def test_dispatch_to_bark(self):
+        raw = json.dumps({"type": "bark", "agent_id": "t", "content": "woof", "ts": 1.0})
+        msg = WSMessage.from_json(raw)
+        assert isinstance(msg, BarkMessage)
+        assert msg.content == "woof"
+
+    def test_dispatch_to_error(self):
+        raw = json.dumps({"type": "error", "code": "E1", "message": "oops", "ts": 1.0})
+        msg = WSMessage.from_json(raw)
+        assert isinstance(msg, ErrorMessage)
+
+    def test_unknown_type_returns_base(self):
+        raw = json.dumps({"type": "unknown_future_type", "ts": 1.0})
+        msg = WSMessage.from_json(raw)
+        assert type(msg) is WSMessage
+        assert msg.type == "unknown_future_type"
+
+    def test_invalid_json_raises(self):
+        with pytest.raises(json.JSONDecodeError):
+            WSMessage.from_json("not json")