Timmy-time-dashboard/tests/unit/test_protocol.py

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