147 lines
4.6 KiB
Python
147 lines
4.6 KiB
Python
"""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
|