183 lines
7.2 KiB
Python
183 lines
7.2 KiB
Python
"""Tests for cloud/memory.py (MemoryStore) and cloud/kb.py (KBStore).
|
||
|
||
Run from ~/workspace/codex-app-server:
|
||
python3 -m unittest discover -s cloud/tests -p "test_memory_kb.py" -v
|
||
|
||
No network, no real model turns, no secrets. Each test uses a fresh
|
||
temporary SQLite file.
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import tempfile
|
||
import unittest
|
||
|
||
sys.path.insert(0, os.path.normpath(os.path.join(os.path.dirname(__file__), "..")))
|
||
|
||
from kb import KBStore # noqa: E402
|
||
from memory import MemoryStore # noqa: E402
|
||
|
||
|
||
def _tmp_db():
|
||
fd, path = tempfile.mkstemp(suffix=".db")
|
||
os.close(fd)
|
||
os.unlink(path) # let sqlite create it fresh
|
||
return path
|
||
|
||
|
||
class MemoryStoreTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.db = _tmp_db()
|
||
self.store = MemoryStore(self.db)
|
||
|
||
def tearDown(self):
|
||
try:
|
||
self.store.conn.close()
|
||
finally:
|
||
if os.path.exists(self.db):
|
||
os.unlink(self.db)
|
||
|
||
def test_add_then_list(self):
|
||
sid = self.store.add_summary(
|
||
"proj-a", "thread-1", "add login page", "login page done",
|
||
files_changed=["app.py"], open_issues=["tests pending"],
|
||
)
|
||
self.assertIsInstance(sid, int)
|
||
rows = self.store.list_summaries("proj-a")
|
||
self.assertEqual(len(rows), 1)
|
||
self.assertEqual(rows[0]["thread_id"], "thread-1")
|
||
self.assertEqual(rows[0]["goal"], "add login page")
|
||
self.assertEqual(rows[0]["files_changed"], ["app.py"])
|
||
self.assertEqual(rows[0]["open_issues"], ["tests pending"])
|
||
self.assertEqual(rows[0]["status"], "active")
|
||
|
||
def test_deprecate_hidden_by_default_visible_with_flag(self):
|
||
sid = self.store.add_summary("proj-a", "t1", "g1", "c1")
|
||
self.store.add_summary("proj-a", "t2", "g2", "c2")
|
||
self.store.deprecate(sid)
|
||
default_rows = self.store.list_summaries("proj-a")
|
||
self.assertEqual([r["thread_id"] for r in default_rows], ["t2"])
|
||
all_rows = self.store.list_summaries("proj-a", include_deprecated=True)
|
||
self.assertEqual(len(all_rows), 2)
|
||
deprecated = [r for r in all_rows if r["id"] == sid]
|
||
self.assertEqual(len(deprecated), 1)
|
||
self.assertEqual(deprecated[0]["status"], "deprecated")
|
||
|
||
def test_delete_removes(self):
|
||
sid = self.store.add_summary("proj-a", "t1", "g1", "c1")
|
||
self.store.delete(sid)
|
||
self.assertEqual(self.store.list_summaries("proj-a"), [])
|
||
self.assertEqual(
|
||
self.store.list_summaries("proj-a", include_deprecated=True), []
|
||
)
|
||
|
||
def test_build_prompt_context_top_k_and_excludes_deprecated(self):
|
||
# 7 active summaries; only the newest 5 may appear.
|
||
for i in range(7):
|
||
self.store.add_summary("proj-a", "t%d" % i, "goal-%d" % i, "done-%d" % i)
|
||
ctx = self.store.build_prompt_context("proj-a", top_k=5)
|
||
self.assertEqual(ctx.count("\n") + 1, 5)
|
||
for i in range(2, 7):
|
||
self.assertIn("goal-%d" % i, ctx)
|
||
for i in range(2):
|
||
self.assertNotIn("goal-%d" % i, ctx)
|
||
# Format spot-check on one line.
|
||
self.assertIn("- [t6] 目标：goal-6；结论：done-6；", ctx)
|
||
# A deprecated newest summary must not appear even within top_k.
|
||
sid = self.store.add_summary("proj-a", "t-new", "goal-new", "done-new")
|
||
self.store.deprecate(sid)
|
||
ctx2 = self.store.build_prompt_context("proj-a", top_k=5)
|
||
self.assertNotIn("goal-new", ctx2)
|
||
self.assertIn("goal-6", ctx2)
|
||
|
||
def test_build_prompt_context_empty(self):
|
||
self.assertEqual(self.store.build_prompt_context("proj-empty"), "")
|
||
|
||
|
||
class KBStoreTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.db = _tmp_db()
|
||
self.store = KBStore(self.db)
|
||
|
||
def tearDown(self):
|
||
try:
|
||
self.store.conn.close()
|
||
finally:
|
||
if os.path.exists(self.db):
|
||
os.unlink(self.db)
|
||
|
||
def test_search_finds_file_by_keyword(self):
|
||
self.store.index_file(
|
||
"proj-a", "docs/auth.md",
|
||
"authentication flow uses token refresh and session cookies",
|
||
kind="doc",
|
||
)
|
||
self.store.index_file(
|
||
"proj-a", "docs/other.md", "unrelated gardening notes", kind="doc"
|
||
)
|
||
results = self.store.search("proj-a", "authentication")
|
||
self.assertEqual(len(results), 1)
|
||
self.assertEqual(results[0]["path"], "docs/auth.md")
|
||
self.assertEqual(results[0]["kind"], "doc")
|
||
self.assertIn("authentication", results[0]["snippet"])
|
||
self.assertLessEqual(len(results[0]["snippet"]), 200)
|
||
|
||
def test_search_project_isolation(self):
|
||
self.store.index_file(
|
||
"proj-a", "docs/secret.md", "zebra secret keyword", kind="doc"
|
||
)
|
||
self.store.index_file(
|
||
"proj-b", "docs/plain.md", "plain notes with no keyword", kind="doc"
|
||
)
|
||
# A's doc must never leak into B's search, and vice versa.
|
||
self.assertEqual(self.store.search("proj-b", "zebra"), [])
|
||
self.assertEqual(self.store.search("proj-a", "zebra")[0]["path"], "docs/secret.md")
|
||
|
||
def test_reindex_same_path_no_duplicates(self):
|
||
self.store.index_file("proj-a", "README.md", "version one content", kind="readme")
|
||
self.store.index_file("proj-a", "README.md", "version two content", kind="readme")
|
||
docs = self.store.list_docs("proj-a")
|
||
self.assertEqual(len(docs), 1)
|
||
self.assertEqual(docs[0]["path"], "README.md")
|
||
self.assertEqual(docs[0]["content"], "version two content")
|
||
# Search must also see only the new version, once.
|
||
self.assertEqual(len(self.store.search("proj-a", "version")), 1)
|
||
self.assertEqual(self.store.search("proj-a", "version")[0]["snippet"],
|
||
"version two content")
|
||
self.assertEqual(self.store.search("proj-a", "one"), [])
|
||
|
||
def test_index_task_and_list_docs(self):
|
||
tid = self.store.index_task(
|
||
"proj-a", "task-42", "fix the login bug", "fixed by resetting tokens"
|
||
)
|
||
self.assertIsInstance(tid, int)
|
||
docs = self.store.list_docs("proj-a")
|
||
self.assertEqual(len(docs), 1)
|
||
self.assertEqual(docs[0]["path"], "task:task-42")
|
||
self.assertEqual(docs[0]["kind"], "task")
|
||
results = self.store.search("proj-a", "login")
|
||
self.assertEqual(len(results), 1)
|
||
self.assertEqual(results[0]["path"], "task:task-42")
|
||
|
||
def test_search_chinese_keyword(self):
|
||
# FTS5 unicode61 does not tokenize CJK words, so a verbatim
|
||
# Chinese keyword must still be found via the LIKE fallback.
|
||
self.store.index_file(
|
||
"proj-a", "docs/login.md", "用户登录流程使用令牌刷新", kind="doc"
|
||
)
|
||
results = self.store.search("proj-a", "登录")
|
||
self.assertEqual(len(results), 1)
|
||
self.assertEqual(results[0]["path"], "docs/login.md")
|
||
|
||
def test_search_top_k_and_empty_query(self):
|
||
for i in range(4):
|
||
self.store.index_file(
|
||
"proj-a", "doc%d.md" % i, "shared keyword doc %d" % i, kind="doc"
|
||
)
|
||
self.assertEqual(len(self.store.search("proj-a", "shared", top_k=2)), 2)
|
||
self.assertEqual(self.store.search("proj-a", ""), [])
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|