1 files changed, 175 insertions, 0 deletions
diff --git a/test_llamachat.py b/test_llamachat.py
index 4a113e0..7707b7b 100755
--- a/test_llamachat.py
+++ b/test_llamachat.py
@@ -14,6 +14,7 @@
 # GNU General Public License for more details.
 """Self-checks for the non-GUI logic. Run: ./test_llamachat.py"""
 
+import json
 import os
 import sys
 import tempfile
@@ -409,6 +410,175 @@ def test_system_qt_theme_guard():
     print("ok  system qt theme guard")
 
 
+def test_prompt_store():
+    """Prompts are files; a preset replaces the global one."""
+    from llamachat import prompts
+
+    with tempfile.TemporaryDirectory() as tmp:
+        store = prompts.PromptStore(Path(tmp) / "prompts")
+        assert store.names() == []
+
+        store.ensure_default()
+        assert store.names() == ["default"]
+        assert store.path_for("default").exists()
+
+        # ensure_default must not clobber an edited global prompt.
+        store.save("default", "edited global")
+        store.ensure_default()
+        assert store.load("default").text == "edited global"
+
+        store.save("coding", "you write code")
+        store.save("terse", "be brief")
+        # Global first, then presets alphabetically.
+        assert store.names() == ["default", "coding", "terse"]
+
+        # A preset replaces the global prompt rather than adding to it.
+        assert store.resolve("coding") == "you write code"
+        assert store.resolve("") == "edited global"
+        assert store.resolve("default") == "edited global"
+
+        # Sentinels.
+        assert store.resolve(prompts.NONE) == ""
+        assert store.resolve(prompts.CUSTOM, "one off") == "one off"
+
+        # A preset deleted behind our back falls back rather than sending
+        # nothing, which would silently change the model's behaviour.
+        assert store.resolve("missing") == "edited global"
+
+        assert store.delete("coding") is True
+        assert store.names() == ["default", "terse"]
+        # The global prompt is the fallback for everything, so it stays.
+        assert store.delete("default") is False
+        assert "default" in store.names()
+
+        # Names become filenames, so path separators must not escape.
+        assert prompts.safe_name("../../etc/passwd") == "etc-passwd"
+        assert prompts.safe_name("/etc/passwd") == "etc-passwd"
+        assert prompts.safe_name("my prompt!") == "my-prompt"
+        for hostile in ("../escape", "a/b", "/tmp/x"):
+            saved = store.save(hostile, "x")
+            assert saved.path.parent == store.dir, saved.path
+            assert saved.path.resolve().parent == store.dir.resolve()
+
+        # A name that reduces to nothing is rejected rather than writing
+        # a dotfile called ".md".
+        for empty in ("..", "", "....", "///"):
+            assert prompts.safe_name(empty) == ""
+            assert store.load(empty) is None
+            assert store.delete(empty) is False
+            try:
+                store.path_for(empty)
+            except ValueError:
+                pass
+            else:
+                raise AssertionError(f"{empty!r} should be rejected")
+    print("ok  prompt store")
+
+
+def test_token_estimate():
+    """The meter's pre-send estimate tracks message size."""
+    small = [{"role": "user", "content": "hi"}]
+    large = [{"role": "user", "content": "word " * 1000}]
+
+    assert backend.estimate_tokens([], 3.5) == 0
+    assert backend.estimate_tokens(small, 3.5) < 20
+    assert backend.estimate_tokens(large, 3.5) > 1000
+
+    # More history means a bigger prompt.
+    grown = large + [{"role": "assistant", "content": "reply " * 500}]
+    assert backend.estimate_tokens(grown, 3.5) > backend.estimate_tokens(
+        large, 3.5
+    )
+
+    # An image costs far more than its text part suggests.
+    with_image = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "what is this"},
+                {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA"}},
+            ],
+        }
+    ]
+    assert backend.estimate_tokens(with_image, 3.5) > 500
+
+    # A silly ratio must not divide by zero.
+    assert backend.estimate_tokens(small, 0) > 0
+    print("ok  token estimate")
+
+
+def test_usage_parsing():
+    """The final stream chunk carries the exact prompt token count."""
+    line = (
+        'data: {"choices":[],"usage":{"prompt_tokens":1234,'
+        '"completion_tokens":56,"total_tokens":1290}}'
+    )
+    kind, payload = backend._parse_sse_line(line)
+    assert kind == "usage"
+    stats = json.loads(payload)
+    assert stats["prompt_tokens"] == 1234
+    assert stats["total_tokens"] == 1290
+
+    # A usage-less chunk with empty choices is not mistaken for one.
+    assert backend._parse_sse_line('data: {"choices":[]}') is None
+    print("ok  usage parsing")
+
+
+def test_session_prompt_storage():
+    """A conversation remembers the prompt it was built with."""
+    from llamachat import prompts
+
+    with tempfile.TemporaryDirectory() as tmp:
+        history = db.History(Path(tmp) / "p.db")
+
+        sid = history.create_session(
+            "chat", "m", "t", prompt_name="coding", prompt_custom=""
+        )
+        row = history.get_session(sid)
+        assert row["prompt_name"] == "coding"
+        assert row["prompt_custom"] == ""
+
+        history.set_prompt(sid, prompts.CUSTOM, "just this once")
+        row = history.get_session(sid)
+        assert row["prompt_name"] == prompts.CUSTOM
+        assert row["prompt_custom"] == "just this once"
+
+        # Defaults keep older call sites working.
+        plain = history.create_session("oneshot", "m", "t")
+        assert history.get_session(plain)["prompt_name"] == ""
+        history.close()
+    print("ok  session prompt storage")
+
+
+def test_prompt_column_migration():
+    """A database predating the prompt columns must still open."""
+    import sqlite3
+
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "old.db"
+        conn = sqlite3.connect(path)
+        conn.executescript(
+            "CREATE TABLE sessions (id INTEGER PRIMARY KEY, mode TEXT,"
+            " title TEXT, model TEXT, created_at INTEGER, updated_at INTEGER);"
+            "CREATE TABLE messages (id INTEGER PRIMARY KEY, session_id INTEGER,"
+            " role TEXT NOT NULL, content TEXT NOT NULL,"
+            " created_at INTEGER NOT NULL);"
+            "INSERT INTO sessions VALUES (1,'chat','old','m',0,0);"
+        )
+        conn.commit()
+        conn.close()
+
+        history = db.History(path)
+        row = history.get_session(1)
+        assert row["title"] == "old"
+        assert row["prompt_name"] == ""
+        assert row["prompt_custom"] == ""
+        history.set_prompt(1, "coding")
+        assert history.get_session(1)["prompt_name"] == "coding"
+        history.close()
+    print("ok  prompt column migration")
+
+
 def test_version_matches_changelog():
     """The package version must be the newest release in the changelog."""
     import re
@@ -506,6 +676,11 @@ if __name__ == "__main__":
     test_reasoning_storage()
     test_migration_adds_reasoning()
     test_user_content()
+    test_prompt_store()
+    test_token_estimate()
+    test_usage_parsing()
+    test_session_prompt_storage()
+    test_prompt_column_migration()
     test_markdown_rendering()
     test_system_qt_theme_guard()
     test_version_matches_changelog()