59 lines
2.0 KiB
Python
59 lines
2.0 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from javis.memory.sqlite_store import SessionNotFoundError, SQLiteSessionStore
|
|
|
|
|
|
class SQLiteSessionStoreTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.temporary_directory = tempfile.TemporaryDirectory()
|
|
database = Path(self.temporary_directory.name) / "sessions.sqlite3"
|
|
self.store = SQLiteSessionStore(database)
|
|
|
|
def tearDown(self) -> None:
|
|
self.temporary_directory.cleanup()
|
|
|
|
def test_create_list_and_load_session(self) -> None:
|
|
created = self.store.create_session("ollama", "test-model")
|
|
loaded = self.store.get_session(created.id)
|
|
|
|
self.assertEqual(loaded.id, created.id)
|
|
self.assertEqual(loaded.provider, "ollama")
|
|
self.assertEqual(loaded.model, "test-model")
|
|
self.assertEqual(self.store.list_sessions()[0].id, created.id)
|
|
|
|
def test_messages_keep_exchange_order(self) -> None:
|
|
session = self.store.create_session("ollama", "test-model")
|
|
self.store.append_exchange(session.id, "eins", "zwei")
|
|
self.store.append_exchange(session.id, "drei", "vier")
|
|
|
|
messages = self.store.get_messages(session.id)
|
|
|
|
self.assertEqual(
|
|
[(message.role, message.content) for message in messages],
|
|
[
|
|
("user", "eins"),
|
|
("assistant", "zwei"),
|
|
("user", "drei"),
|
|
("assistant", "vier"),
|
|
],
|
|
)
|
|
|
|
def test_unknown_session_raises_clear_error(self) -> None:
|
|
with self.assertRaises(SessionNotFoundError):
|
|
self.store.get_session("missing")
|
|
|
|
def test_clear_keeps_session_and_removes_messages(self) -> None:
|
|
session = self.store.create_session("ollama", "test-model")
|
|
self.store.append_exchange(session.id, "eins", "zwei")
|
|
|
|
self.store.clear_messages(session.id)
|
|
|
|
self.assertEqual(self.store.get_messages(session.id), [])
|
|
self.assertEqual(self.store.get_session(session.id).message_count, 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|