import asyncio import hashlib import json import os import sqlite3 import sys import tempfile import threading import types import unittest from pathlib import Path from unittest import mock BASE = Path(__file__).resolve().parents[1] sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)] # These tests never perform HTTP. Keep them offline on minimal Python installs. try: import aiohttp # noqa: F401 except ModuleNotFoundError: aiohttp_stub = types.ModuleType("aiohttp") class OfflineClientSession: def __init__(self, *args, **kwargs): raise RuntimeError("HTTP is disabled in offline chat migration tests") aiohttp_stub.ClientSession = OfflineClientSession sys.modules["aiohttp"] = aiohttp_stub from chat import engine import migrate_chat_data DEFAULT_ENGINE_PATHS = { "history": engine.HISTORY_DIR, "persona": engine.PROMPT_FILE, "lorebook": engine.LOREBOOK_PATH, } class RecordingChannel: def __init__(self, name="tg"): self.name = name self.sent = [] async def send(self, msg_type, target_id, content): self.sent.append((msg_type, target_id, content)) return True class EngineTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory(prefix="chat-engine-test-") root = Path(self.tmp.name) self.paths = { "BASE_DIR": str(root), "HISTORY_DIR": str(root / "history"), "GROUP_HISTORY_DIR": str(root / "history" / "group"), "MEMORY_DIR": str(root / "memory"), "STATUS_DIR": str(root / "status"), "SETTING_FILE": str(root / "setting.json"), "STICKERS_DIR": str(root / "stickers"), "PROMPTS_DIR": str(root / "prompts"), "LOREBOOK_PATH": str(root / "prompts" / "lorebook.json"), } self.path_patch = mock.patch.multiple(engine, **self.paths) self.path_patch.start() for key in ("HISTORY_DIR", "GROUP_HISTORY_DIR", "MEMORY_DIR", "STATUS_DIR", "STICKERS_DIR", "PROMPTS_DIR"): os.makedirs(self.paths[key], exist_ok=True) engine._delivery_targets.clear() engine._user_channel.clear() engine._history_revisions.clear() engine._websocket = None def tearDown(self): engine._delivery_targets.clear() engine._user_channel.clear() engine._history_revisions.clear() engine._websocket = None self.path_patch.stop() self.tmp.cleanup() def test_package_paths_and_prompt_defaults(self): package = Path(engine.__file__).resolve().parent self.assertEqual(Path(DEFAULT_ENGINE_PATHS["history"]), package / "history") self.assertEqual(Path(DEFAULT_ENGINE_PATHS["persona"]), package.parent / "prompts" / "persona.txt") self.assertEqual(Path(DEFAULT_ENGINE_PATHS["lorebook"]), package.parent / "prompts" / "lorebook.json") # The temporary prompt directory is empty: built-in defaults must work. self.assertIn("JSON", engine.load_instruction("agent_instructions.txt")) self.assertIn("群聊", engine.load_instruction("group_instructions.txt")) Path(engine.PROMPTS_DIR, "compact_instructions.txt").write_text("自定义压缩指令", encoding="utf-8") self.assertEqual(engine.load_instruction("compact_instructions.txt"), "自定义压缩指令") def test_private_model_request_has_no_legacy_debounce(self): self.assertEqual(engine.USER_IDLE_WAIT_MIN, 0) self.assertEqual(engine.USER_IDLE_WAIT_MAX, 0) def test_explicit_delivery_never_uses_uid_as_target(self): uid = 9000000001 channel = RecordingChannel() engine.set_delivery(uid, channel, 10000000042) ok = asyncio.run(engine.send_private_message(None, uid, "你好。")) self.assertTrue(ok) self.assertEqual(channel.sent, [("private", 10000000042, "你好")]) self.assertNotEqual(channel.sent[0][1], uid) def test_restart_recovers_one_identity_target_without_mirror_fanout(self): uid = 9000000001 channel = RecordingChannel() engine.set_websocket(channel) # only channel registry fallback; no delivery binding survives with mock.patch("identity.push_target", return_value=("tg", 10000000042)) as preferred, \ mock.patch("identity.push_targets", return_value=[("tg", 10000000042), ("qq", 12345)]) as fanout: ok = asyncio.run(engine.send_private_message(None, uid, "单条投递")) self.assertTrue(ok) preferred.assert_called_once_with(uid) fanout.assert_not_called() self.assertEqual(channel.sent, [("private", 10000000042, "单条投递")]) def test_qq_only_identity_uses_first_target_only(self): uid = 9000000001 channel = RecordingChannel("qq") engine.set_websocket(channel) with mock.patch("identity.push_target", return_value=None), \ mock.patch("identity.push_targets", return_value=[("qq", 12345), ("qq", 67890)]): asyncio.run(engine.send_private_message(None, uid, "QQ")) self.assertEqual(channel.sent, [("private", 12345, "QQ")]) def test_unmapped_uid_fails_closed(self): channel = RecordingChannel() with mock.patch("identity.push_target", return_value=None), \ mock.patch("identity.push_targets", return_value=[]): ok = asyncio.run(engine.send_private_message(channel, 9000000001, "不应发送")) self.assertFalse(ok) self.assertEqual(channel.sent, []) def test_compact_cas_does_not_drop_concurrent_append(self): uid = 9000000001 path = Path(engine.HISTORY_DIR) / (str(uid) + ".txt") original = ["[2026-01-01 00:00:0%d][user] line-%d" % (i, i) for i in range(5)] path.write_text("\n".join(original) + "\n", encoding="utf-8") async def model_with_append(*args, **kwargs): engine.append_history(uid, "user", "concurrent-new-line") return "摘要" with mock.patch.object(engine, "call_deepseek", side_effect=model_with_append): result = asyncio.run(engine.compact_history(uid, keep_recent=2)) self.assertIn("未改写历史", result) lines = path.read_text(encoding="utf-8").splitlines() self.assertEqual(lines[:5], original) self.assertIn("concurrent-new-line", lines[-1]) self.assertEqual(len(lines), 6) async def model(*args, **kwargs): return "摘要" with mock.patch.object(engine, "call_deepseek", side_effect=model): result = asyncio.run(engine.compact_history(uid, keep_recent=2)) self.assertIn("已完成", result) lines = path.read_text(encoding="utf-8").splitlines() self.assertIn("[summary]", lines[0]) self.assertEqual(len(lines), 3) self.assertIn("concurrent-new-line", lines[-1]) self.assertEqual(list(Path(engine.HISTORY_DIR).glob(".chat-*.tmp")), []) def test_status_read_modify_write_uses_uid_lock(self): uid = 9000000001 barrier = threading.Barrier(9) def worker(index): barrier.wait() if index % 2: engine.update_user_status(uid, pause=True) else: engine.update_user_status(uid, balance_notice_date="2026-10-29") threads = [threading.Thread(target=worker, args=(i,)) for i in range(8)] for thread in threads: thread.start() barrier.wait() for thread in threads: thread.join() status = engine.load_user_status(uid) self.assertTrue(status["pause"]) self.assertEqual(status["balance_notice_date"], "2026-10-29") def test_modelscope_counter_does_not_lose_whole_file_updates(self): workers = [threading.Thread(target=engine._modelscope_mark_success) for _ in range(12)] for thread in workers: thread.start() for thread in workers: thread.join() self.assertEqual(engine.load_setting()["modelscope_count"], 12) def test_group_prompt_and_reply_never_read_private_memory(self): response = json.dumps({"message": [{"content": "群回复"}]}, ensure_ascii=False) with mock.patch.object(engine, "call_deepseek", new=mock.AsyncMock(return_value=response)) as model, \ mock.patch.object(engine, "send_group_message", new=mock.AsyncMock()), \ mock.patch.object(engine.memory_store, "migrate_legacy") as migrate, \ mock.patch.object(engine.memory_store, "recall_facts") as recall, \ mock.patch.object(engine.memory_store, "get_relation") as relation, \ mock.patch.object(engine.memory_store, "list_entities") as entities: asyncio.run(engine._call_group_model_and_reply_impl(None, 1000000000999, "大家好")) migrate.assert_not_called() recall.assert_not_called() relation.assert_not_called() entities.assert_not_called() prompt = model.await_args.args[0] self.assertIn("不要读取、推断或泄露私聊", prompt) class MigrationTestCase(unittest.TestCase): UID = 9000000001 TG_NS = 10000000042 QQ = 12345 GROUP = 1000000000999 def setUp(self): self.tmp = tempfile.TemporaryDirectory(prefix="chat-migration-test-") root = Path(self.tmp.name) self.source = root / "old-chat" self.accounts = root / "accounts" self.dest = root / "new-chat" self.report = root / "report.json" for directory in (self.source, self.accounts / "users"): directory.mkdir(parents=True, exist_ok=True) record = { "uid": self.UID, "tg": {"raw": 42, "ns": self.TG_NS}, "qq": [{"id": self.QQ}], } (self.accounts / "users" / (str(self.UID) + ".json")).write_text( json.dumps(record), encoding="utf-8") for category in ("history/group", "memory", "status", "stickers"): (self.source / category).mkdir(parents=True, exist_ok=True) (self.source / "history" / (str(self.TG_NS) + ".txt")).write_text( "shared\ntg-only\n", encoding="utf-8") (self.source / "history" / (str(self.QQ) + ".txt")).write_text( "shared\nqq-only\n", encoding="utf-8") (self.source / "history" / "777.txt").write_text("orphan\n", encoding="utf-8") (self.source / "history" / "group" / (str(self.GROUP) + ".txt")).write_text( "group-line\n", encoding="utf-8") (self.source / "memory" / (str(self.TG_NS) + ".txt")).write_text("tg-memory\n", encoding="utf-8") (self.source / "memory" / (str(self.QQ) + ".txt")).write_text("qq-memory\n", encoding="utf-8") (self.source / "status" / (str(self.TG_NS) + ".json")).write_text(json.dumps({ "pause": True, "nested": {"tg": 1}, }), encoding="utf-8") (self.source / "status" / (str(self.QQ) + ".json")).write_text(json.dumps({ "sleep": True, "nested": {"qq": 2}, }), encoding="utf-8") (self.source / "status" / "777.json").write_text("{}", encoding="utf-8") (self.source / "setting.json").write_text(json.dumps({"chat": True}), encoding="utf-8") (self.source / "stickers" / "ok.webp").write_bytes(b"sticker") self._make_source_db() self.before = self._tree_hashes(self.source) def tearDown(self): self.tmp.cleanup() @staticmethod def _tree_hashes(root): return { str(path.relative_to(root)): hashlib.sha256(path.read_bytes()).hexdigest() for path in sorted(root.rglob("*")) if path.is_file() } def _make_source_db(self): con = sqlite3.connect(str(self.source / "memory.db")) con.executescript(migrate_chat_data.SCHEMA) for old_uid, content in ((self.TG_NS, "shared fact"), (self.QQ, "shared fact"), (self.QQ, "qq fact"), (777, "orphan fact")): con.execute( "INSERT INTO facts(uid,content,importance,keywords,created_at,archived,source)" " VALUES(?,?,3,'',1,0,'test')", (old_uid, content)) for old_uid, obj in ((self.TG_NS, "cat"), (self.QQ, "dog"), (777, "orphan")): con.execute( "INSERT INTO triples(uid,subject,predicate,object,importance,valid_from,valid_until,created_at,source)" " VALUES(?,?, 'likes', ?,3,1,NULL,1,'test')", (old_uid, "user", obj)) con.execute("INSERT INTO relations VALUES(?,?,?,?,?,?)", ( self.TG_NS, 10, 2, 50, 20, json.dumps([{"ts": 1, "text": "tg"}]), )) con.execute("INSERT INTO relations VALUES(?,?,?,?,?,?)", ( self.QQ, 5, 3, 60, 30, json.dumps([{"ts": 2, "text": "qq"}]), )) con.execute("INSERT INTO relations VALUES(?,?,?,?,?,?)", (777, 1, 9, 1, 1, "[]")) con.execute("INSERT INTO meta VALUES('imported:%s','1')" % self.TG_NS) con.execute("INSERT INTO meta VALUES('custom','value')") con.commit() con.close() def _counts(self): con = sqlite3.connect(str(self.dest / "memory.db")) result = { table: con.execute("SELECT COUNT(*) FROM %s" % table).fetchone()[0] for table in ("facts", "triples", "relations") } result["relation_count"] = con.execute( "SELECT msg_count FROM relations WHERE uid=?", (self.UID,)).fetchone()[0] result["first_met"] = con.execute( "SELECT first_met FROM relations WHERE uid=?", (self.UID,)).fetchone()[0] result["fts"] = con.execute("SELECT COUNT(*) FROM facts_fts").fetchone()[0] result["fts_hits"] = con.execute( "SELECT COUNT(*) FROM facts_fts WHERE facts_fts MATCH 'shared'").fetchone()[0] con.close() return result def test_full_migration_quarantine_and_idempotency(self): report = migrate_chat_data.migrate( self.source, self.accounts, self.dest, self.report) self.assertTrue(self.report.is_file()) self.assertEqual(report["mapping"]["tg_ns_to_uid"], {str(self.TG_NS): self.UID}) self.assertEqual(report["mapping"]["old_key_to_uid"][str(self.QQ)], self.UID) history = (self.dest / "history" / (str(self.UID) + ".txt")).read_text(encoding="utf-8") self.assertEqual(history.splitlines(), ["shared", "tg-only", "qq-only"]) memory = (self.dest / "memory" / (str(self.UID) + ".txt")).read_text(encoding="utf-8") self.assertEqual(set(memory.splitlines()), {"tg-memory", "qq-memory"}) status = json.loads((self.dest / "status" / (str(self.UID) + ".json")).read_text(encoding="utf-8")) self.assertEqual(status, {"pause": True, "sleep": True, "nested": {"qq": 2, "tg": 1}}) self.assertEqual((self.dest / "history" / "group" / (str(self.GROUP) + ".txt")).read_text( encoding="utf-8"), "group-line\n") self.assertFalse((self.dest / "history" / "777.txt").exists()) self.assertTrue((self.dest / "quarantine" / "history" / "777.txt").is_file()) self.assertTrue((self.dest / "quarantine" / "status" / "777.json").is_file()) self.assertEqual((self.dest / "stickers" / "ok.webp").read_bytes(), b"sticker") self.assertTrue(report["sqlite"]["fts_validated"]) self.assertEqual(report["sqlite"]["unmatched"], {"facts": 1, "triples": 1, "relations": 1}) self.assertEqual(self._counts(), { "facts": 2, "triples": 2, "relations": 1, "relation_count": 5, "first_met": 5, "fts": 2, "fts_hits": 1, }) self.assertEqual(self._tree_hashes(self.source), self.before, "source was modified") # A second run must not duplicate lines/facts/triples or add relation counts again. migrate_chat_data.migrate(self.source, self.accounts, self.dest, self.report) self.assertEqual(self._counts(), { "facts": 2, "triples": 2, "relations": 1, "relation_count": 5, "first_met": 5, "fts": 2, "fts_hits": 1, }) self.assertEqual((self.dest / "history" / (str(self.UID) + ".txt")).read_text( encoding="utf-8"), history) self.assertEqual(self._tree_hashes(self.source), self.before, "source changed on rerun") def test_existing_destination_values_win_json_conflicts(self): (self.dest / "status").mkdir(parents=True) (self.dest / "status" / (str(self.UID) + ".json")).write_text(json.dumps({ "pause": False, "custom": "new", }), encoding="utf-8") report = migrate_chat_data.migrate(self.source, self.accounts, self.dest, self.report) status = json.loads((self.dest / "status" / (str(self.UID) + ".json")).read_text(encoding="utf-8")) self.assertFalse(status["pause"]) self.assertEqual(status["custom"], "new") self.assertTrue(status["sleep"]) self.assertTrue(any(item["kind"] == "json_value" for item in report["conflicts"])) class RawTelegramMigrationTest(unittest.TestCase): def test_explicit_raw_tg_provisioning_avoids_qq_guess(self): with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) source, accounts, dest = root / "source", root / "accounts", root / "dest" (source / "history").mkdir(parents=True) accounts.mkdir(parents=True) (source / "history" / "123.txt").write_text("legacy\n", encoding="utf-8") con = sqlite3.connect(str(source / "memory.db")) con.executescript(migrate_chat_data.SCHEMA) con.execute("INSERT INTO facts(uid,content,importance,keywords,created_at,archived,source) VALUES(123,'fact',3,'',1,0,'test')") con.commit() con.close() report = migrate_chat_data.migrate( source, accounts, dest, raw_tg=True, provision_tg=True) self.assertEqual(len(report["provisioned_tg"]), 1) uid = report["provisioned_tg"][0]["uid"] self.assertEqual(report["provisioned_tg"][0]["tg_ns"], 10000000123) self.assertTrue((dest / "history" / (str(uid) + ".txt")).is_file()) self.assertEqual(report["sqlite"]["unmatched"]["facts"], 0) self.assertTrue((accounts / "users" / (str(uid) + ".json")).is_file()) if __name__ == "__main__": unittest.main()