"""Tests for Session and Message data models.""" import pytest from agentkit.session.models import Message, MessageRole, Session, SessionStatus class TestSessionStatus: def test_status_values(self): assert SessionStatus.ACTIVE == "active" assert SessionStatus.PAUSED == "paused" assert SessionStatus.CLOSED == "closed" def test_status_from_string(self): assert SessionStatus("active") == SessionStatus.ACTIVE assert SessionStatus("paused") == SessionStatus.PAUSED assert SessionStatus("closed") == SessionStatus.CLOSED class TestMessageRole: def test_role_values(self): assert MessageRole.SYSTEM == "system" assert MessageRole.USER == "user" assert MessageRole.ASSISTANT == "assistant" assert MessageRole.TOOL == "tool" class TestSession: def test_create_session(self): session = Session(session_id="s1", agent_name="test-agent") assert session.session_id == "s1" assert session.agent_name == "test-agent" assert session.status == SessionStatus.ACTIVE assert session.metadata == {} assert session.created_at is not None assert session.updated_at is not None def test_session_to_dict_and_back(self): session = Session( session_id="s1", agent_name="agent1", status=SessionStatus.PAUSED, metadata={"key": "value"}, ) d = session.to_dict() assert d["session_id"] == "s1" assert d["agent_name"] == "agent1" assert d["status"] == "paused" assert d["metadata"] == {"key": "value"} restored = Session.from_dict(d) assert restored.session_id == session.session_id assert restored.agent_name == session.agent_name assert restored.status == session.status assert restored.metadata == session.metadata def test_new_session_id_is_unique(self): ids = {Session.new_session_id() for _ in range(100)} assert len(ids) == 100 def test_new_message_id_is_unique(self): ids = {Session.new_message_id() for _ in range(100)} assert len(ids) == 100 class TestMessage: def test_create_message(self): msg = Message( message_id="m1", session_id="s1", role=MessageRole.USER, content="Hello", ) assert msg.message_id == "m1" assert msg.session_id == "s1" assert msg.role == MessageRole.USER assert msg.content == "Hello" assert msg.tool_call_id is None assert msg.agent_name is None assert msg.metadata == {} def test_message_with_tool_call(self): msg = Message( message_id="m1", session_id="s1", role=MessageRole.TOOL, content="result", tool_call_id="tc1", agent_name="agent1", ) assert msg.tool_call_id == "tc1" assert msg.agent_name == "agent1" def test_message_to_dict_and_back(self): msg = Message( message_id="m1", session_id="s1", role=MessageRole.ASSISTANT, content="Hi there", tool_call_id="tc1", agent_name="agent1", metadata={"step": 1}, ) d = msg.to_dict() assert d["message_id"] == "m1" assert d["role"] == "assistant" assert d["tool_call_id"] == "tc1" restored = Message.from_dict(d) assert restored.message_id == msg.message_id assert restored.role == msg.role assert restored.content == msg.content assert restored.tool_call_id == msg.tool_call_id assert restored.agent_name == msg.agent_name assert restored.metadata == msg.metadata def test_to_chat_message_user(self): msg = Message( message_id="m1", session_id="s1", role=MessageRole.USER, content="Hello", ) chat_msg = msg.to_chat_message() assert chat_msg == {"role": "user", "content": "Hello"} def test_to_chat_message_tool(self): msg = Message( message_id="m1", session_id="s1", role=MessageRole.TOOL, content="result", tool_call_id="tc1", ) chat_msg = msg.to_chat_message() assert chat_msg == {"role": "tool", "content": "result", "tool_call_id": "tc1"} def test_to_chat_message_no_tool_call_id(self): msg = Message( message_id="m1", session_id="s1", role=MessageRole.ASSISTANT, content="Hi", ) chat_msg = msg.to_chat_message() assert "tool_call_id" not in chat_msg