158 lines
5.3 KiB
Python
158 lines
5.3 KiB
Python
"""Tests for InMemorySessionStore."""
|
|
|
|
import pytest
|
|
|
|
from agentkit.session.models import Message, MessageRole, Session, SessionStatus
|
|
from agentkit.session.store import InMemorySessionStore
|
|
|
|
|
|
@pytest.fixture
|
|
def store():
|
|
return InMemorySessionStore()
|
|
|
|
|
|
async def _create_session(store, session_id="s1", agent_name="agent1"):
|
|
session = Session(session_id=session_id, agent_name=agent_name)
|
|
await store.save_session(session)
|
|
return session
|
|
|
|
|
|
async def _create_message(store, session_id, role=MessageRole.USER, content="Hello"):
|
|
msg = Message(
|
|
message_id=Session.new_message_id(),
|
|
session_id=session_id,
|
|
role=role,
|
|
content=content,
|
|
)
|
|
await store.append_message(msg)
|
|
return msg
|
|
|
|
|
|
class TestInMemorySessionStoreCRUD:
|
|
@pytest.mark.asyncio
|
|
async def test_save_and_get(self, store):
|
|
session = await _create_session(store)
|
|
fetched = await store.get_session("s1")
|
|
assert fetched is not None
|
|
assert fetched.session_id == "s1"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_nonexistent(self, store):
|
|
result = await store.get_session("nonexistent")
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_status(self, store):
|
|
await _create_session(store)
|
|
updated = await store.update_session_status("s1", SessionStatus.PAUSED)
|
|
assert updated is not None
|
|
assert updated.status == SessionStatus.PAUSED
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_status_nonexistent(self, store):
|
|
result = await store.update_session_status("nonexistent", SessionStatus.PAUSED)
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete(self, store):
|
|
await _create_session(store)
|
|
assert await store.delete_session("s1") is True
|
|
assert await store.get_session("s1") is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_nonexistent(self, store):
|
|
assert await store.delete_session("nonexistent") is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_sessions(self, store):
|
|
await _create_session(store, "s1", "agent1")
|
|
await _create_session(store, "s2", "agent2")
|
|
sessions = await store.list_sessions()
|
|
assert len(sessions) == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_sessions_by_agent(self, store):
|
|
await _create_session(store, "s1", "agent1")
|
|
await _create_session(store, "s2", "agent2")
|
|
sessions = await store.list_sessions(agent_name="agent1")
|
|
assert len(sessions) == 1
|
|
assert sessions[0].agent_name == "agent1"
|
|
|
|
|
|
class TestInMemorySessionStoreMessages:
|
|
@pytest.mark.asyncio
|
|
async def test_append_and_get(self, store):
|
|
await _create_session(store)
|
|
await _create_message(store, "s1", content="Hello")
|
|
await _create_message(store, "s1", content="World")
|
|
|
|
messages = await store.get_messages("s1")
|
|
assert len(messages) == 2
|
|
assert messages[0].content == "Hello"
|
|
assert messages[1].content == "World"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_messages_pagination(self, store):
|
|
await _create_session(store)
|
|
for i in range(5):
|
|
await _create_message(store, "s1", content=f"Msg {i}")
|
|
|
|
page = await store.get_messages("s1", limit=2, offset=1)
|
|
assert len(page) == 2
|
|
assert page[0].content == "Msg 1"
|
|
assert page[1].content == "Msg 2"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_count_messages(self, store):
|
|
await _create_session(store)
|
|
await _create_message(store, "s1")
|
|
await _create_message(store, "s1")
|
|
assert await store.count_messages("s1") == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_count_messages_empty_session(self, store):
|
|
assert await store.count_messages("nonexistent") == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_messages_empty_session(self, store):
|
|
messages = await store.get_messages("nonexistent")
|
|
assert messages == []
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_session_removes_messages(self, store):
|
|
await _create_session(store)
|
|
await _create_message(store, "s1")
|
|
await store.delete_session("s1")
|
|
assert await store.count_messages("s1") == 0
|
|
|
|
|
|
class TestInMemorySessionStoreEviction:
|
|
@pytest.mark.asyncio
|
|
async def test_evict_closed_session_on_full(self):
|
|
store = InMemorySessionStore(max_sessions=2)
|
|
s1 = await _create_session(store, "s1")
|
|
await _create_session(store, "s2")
|
|
|
|
# Close s1 so it can be evicted
|
|
await store.update_session_status("s1", SessionStatus.CLOSED)
|
|
|
|
# Creating a third session should evict s1
|
|
await _create_session(store, "s3")
|
|
assert await store.get_session("s1") is None
|
|
assert await store.get_session("s2") is not None
|
|
assert await store.get_session("s3") is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_full_no_closed_raises(self):
|
|
store = InMemorySessionStore(max_sessions=2)
|
|
await _create_session(store, "s1")
|
|
await _create_session(store, "s2")
|
|
with pytest.raises(RuntimeError, match="full"):
|
|
await _create_session(store, "s3")
|
|
|
|
|
|
class TestInMemorySessionStoreHealth:
|
|
@pytest.mark.asyncio
|
|
async def test_health_check(self, store):
|
|
assert await store.health_check() is True
|