Initial import of Shirobana bot backend
Include utility modules, shared channels, prompts, read-only skills, wardrobe memory and regression tests. Externalize legacy credentials and exclude private runtime state.
This commit is contained in:
RRRosmontis
@@ -0,0 +1,380 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user