172 lines
6.7 KiB
Python
172 lines
6.7 KiB
Python
import sqlite3
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from javis.memory.sqlite_store import (
|
|
AmbiguousSessionReferenceError,
|
|
SessionNotFoundError,
|
|
SQLiteSessionStore,
|
|
)
|
|
|
|
|
|
class SQLiteSessionStoreTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.temporary_directory = tempfile.TemporaryDirectory()
|
|
self.database = Path(self.temporary_directory.name) / "sessions.sqlite3"
|
|
self.store = SQLiteSessionStore(self.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(loaded.title, "Neue Sitzung")
|
|
self.assertEqual(loaded.last_provider, "ollama")
|
|
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)
|
|
|
|
def test_first_exchange_sets_title_and_last_provider_atomically(self) -> None:
|
|
session = self.store.create_session("hybrid", "auto")
|
|
|
|
self.store.append_exchange(
|
|
session.id,
|
|
"Frage",
|
|
"Antwort",
|
|
title_if_first="Kurzer Titel",
|
|
last_provider="gemini",
|
|
)
|
|
self.store.append_exchange(
|
|
session.id,
|
|
"Noch eine Frage",
|
|
"Noch eine Antwort",
|
|
title_if_first="Darf Titel nicht ersetzen",
|
|
last_provider="ollama",
|
|
)
|
|
|
|
loaded = self.store.get_session(session.id)
|
|
self.assertEqual(loaded.title, "Kurzer Titel")
|
|
self.assertEqual(loaded.last_provider, "ollama")
|
|
|
|
def test_old_database_is_migrated_idempotently_without_data_loss(self) -> None:
|
|
database = Path(self.temporary_directory.name) / "legacy.sqlite3"
|
|
with sqlite3.connect(database) as connection:
|
|
connection.executescript(
|
|
"""
|
|
CREATE TABLE sessions (
|
|
id TEXT PRIMARY KEY,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL,
|
|
provider TEXT NOT NULL,
|
|
model TEXT NOT NULL
|
|
);
|
|
CREATE TABLE messages (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
session_id TEXT NOT NULL,
|
|
position INTEGER NOT NULL,
|
|
role TEXT NOT NULL,
|
|
content TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
UNIQUE (session_id, position)
|
|
);
|
|
INSERT INTO sessions VALUES
|
|
('legacy', '2026-01-01', '2026-01-02', 'ollama', 'old-model');
|
|
INSERT INTO messages
|
|
(session_id, position, role, content, created_at)
|
|
VALUES ('legacy', 0, 'user', 'Alte Nachricht', '2026-01-02');
|
|
"""
|
|
)
|
|
connection.close()
|
|
|
|
migrated = SQLiteSessionStore(database)
|
|
migrated_again = SQLiteSessionStore(database)
|
|
|
|
session = migrated_again.get_session("legacy")
|
|
self.assertEqual(session.title, "Neue Sitzung")
|
|
self.assertEqual(session.last_provider, "ollama")
|
|
self.assertEqual(
|
|
[message.content for message in migrated.get_messages("legacy")],
|
|
["Alte Nachricht"],
|
|
)
|
|
|
|
def test_rename_and_search_only_titles_and_user_messages(self) -> None:
|
|
first = self.store.create_session("ollama", "test-model")
|
|
second = self.store.create_session("ollama", "test-model")
|
|
self.store.rename_session(first.id, "SQLite Hilfe")
|
|
self.store.append_exchange(first.id, "Transaktion erklären", "Privates Lösungswort")
|
|
self.store.append_exchange(second.id, "Andere Frage", "Nur SQLite in Antwort")
|
|
|
|
self.assertEqual(self.store.get_session(first.id).title, "SQLite Hilfe")
|
|
self.assertEqual(
|
|
[session.id for session in self.store.search_sessions("Transaktion")],
|
|
[first.id],
|
|
)
|
|
self.assertEqual(self.store.search_sessions("Lösungswort"), [])
|
|
self.assertEqual(self.store.search_sessions("SQLite"), [self.store.get_session(first.id)])
|
|
|
|
def test_search_treats_sql_wildcards_and_injection_as_plain_text(self) -> None:
|
|
session = self.store.create_session("ollama", "test-model")
|
|
self.store.rename_session(session.id, "100% SQLite_Name")
|
|
|
|
self.assertEqual(self.store.search_sessions("%")[0].id, session.id)
|
|
self.assertEqual(self.store.search_sessions("_")[0].id, session.id)
|
|
self.assertEqual(self.store.search_sessions("' OR 1=1 --"), [])
|
|
|
|
def test_unique_and_ambiguous_session_prefixes(self) -> None:
|
|
with sqlite3.connect(self.database) as connection:
|
|
connection.executemany(
|
|
"""
|
|
INSERT INTO sessions
|
|
(id, created_at, updated_at, provider, model, title, last_provider)
|
|
VALUES (?, '2026-01-01', '2026-01-01', 'ollama', 'model', ?, 'ollama')
|
|
""",
|
|
(
|
|
("abc111", "Eins"),
|
|
("abc222", "Zwei"),
|
|
("unique333", "Drei"),
|
|
),
|
|
)
|
|
connection.close()
|
|
|
|
self.assertEqual(self.store.resolve_unique_prefix("unique").id, "unique333")
|
|
with self.assertRaises(AmbiguousSessionReferenceError):
|
|
self.store.resolve_unique_prefix("abc")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|