337 lines
17 KiB
Python
337 lines
17 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
会话复制（DBManager.copy_session）单测
|
||
—— 全程临时 DB + 临时附件目录，绝不触碰真实 data/chat_history.db
|
||
|
||
覆盖：
|
||
T1 线性会话复制（消息数/顺序/内容一致，ID 全新）
|
||
T2 parent_id 链重建正确（首条 parent 为 None，链走通）
|
||
T3 分叉会话整树复制（分支数一致）
|
||
T4 压缩标记原样复制（数量/内容/切点关系/时间戳晚于全部保留行）
|
||
T5 current_leaf_msg_id 重映射且真实存在
|
||
T6 源会话零改动（全列快照比对）
|
||
T7 删副本 → 源完好；删源 → 副本完好
|
||
T8 图片附件物理复制到新路径（新旧文件同时存在、内容相同）
|
||
T9 文本附件原样保留（不产生新文件、metadata 逐字节相同）
|
||
T10 title / is_starred / mode / has_messages / sort_order 语义
|
||
T11 不存在的 session_id → None，且无残留
|
||
T12 副本再复制 → 标题 (副本 2)
|
||
T13 外键无违规 + 无孤儿 parent_id
|
||
T14 get_message_chain(源) 与 副本 的 (role, content, is_ignored) 序列完全相同
|
||
|
||
运行: PYTHONIOENCODING=utf-8 python tests/test_copy_session.py
|
||
"""
|
||
import os
|
||
import sys
|
||
import json
|
||
import uuid
|
||
import time
|
||
import tempfile
|
||
import shutil
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
from core.db_manager import DBManager # noqa: E402
|
||
|
||
RESULTS = []
|
||
|
||
|
||
def check(name, cond, extra=""):
|
||
RESULTS.append((name, bool(cond)))
|
||
print(f"{'PASS' if cond else 'FAIL'} {name} {extra if not cond else ''}", flush=True)
|
||
|
||
|
||
_TMP = tempfile.mkdtemp(prefix="hocode_copysess_")
|
||
_ATT = os.path.join(_TMP, "data", "attachments")
|
||
os.makedirs(_ATT, exist_ok=True)
|
||
|
||
db = DBManager(os.path.join(_TMP, "data", "chat_history.db"))
|
||
db.files_root = _TMP # 附件根指向临时目录
|
||
|
||
|
||
# ======================================================================
|
||
# 工具函数
|
||
# ======================================================================
|
||
def mk_session(title, msgs, mode=None, starred=0, has_messages=1, att_meta=None):
|
||
"""建一个线性会话。msgs=[(role, content)]；att_meta={index: json_string}
|
||
返回 (session_id, [msg_id])"""
|
||
sid = "sess_" + uuid.uuid4().hex[:12]
|
||
now = int(time.time())
|
||
with db.get_connection() as conn:
|
||
cur = conn.cursor()
|
||
cur.execute(
|
||
"INSERT INTO sessions (id,title,created_at,updated_at,has_messages,sort_order,"
|
||
"is_starred,current_leaf_msg_id,mode) VALUES (?,?,?,?,?,?,?,?,?)",
|
||
(sid, title, now, now, has_messages, db.get_min_sort_order(), starred, None, mode))
|
||
prev, ids = None, []
|
||
for i, (role, content) in enumerate(msgs):
|
||
mid = ("comp_" if role == "compaction" else "msg_") + uuid.uuid4().hex[:16]
|
||
meta = (att_meta or {}).get(i)
|
||
cur.execute(
|
||
"INSERT INTO messages (id,session_id,role,content,reasoning,is_ignored,"
|
||
"created_at,attachment_metadata,parent_id,timeline,usage) "
|
||
"VALUES (?,?,?,?,?,?,?,?,?,?,?)",
|
||
(mid, sid, role, content, "思考内容", 0, now - 1000 + i, meta, prev,
|
||
'{"t":1}', '{"input":10,"output":5}'))
|
||
ids.append(mid)
|
||
prev = mid
|
||
cur.execute("UPDATE sessions SET current_leaf_msg_id = ? WHERE id = ?", (prev, sid))
|
||
conn.commit()
|
||
return sid, ids
|
||
|
||
|
||
def add_child(session_id, parent_id, role, content, ts):
|
||
"""在指定 parent 下挂一个子消息（构造分叉用）"""
|
||
mid = "msg_" + uuid.uuid4().hex[:16]
|
||
with db.get_connection() as conn:
|
||
conn.execute(
|
||
"INSERT INTO messages (id,session_id,role,content,reasoning,is_ignored,"
|
||
"created_at,attachment_metadata,parent_id,timeline,usage) "
|
||
"VALUES (?,?,?,?,?,?,?,?,?,?,?)",
|
||
(mid, session_id, role, content, "", 0, ts, None, parent_id, None, None))
|
||
conn.commit()
|
||
return mid
|
||
|
||
|
||
def rows_of(sid):
|
||
with db.get_connection() as conn:
|
||
return [dict(r) for r in conn.execute(
|
||
"SELECT * FROM messages WHERE session_id = ? ORDER BY created_at ASC, rowid ASC",
|
||
(sid,)).fetchall()]
|
||
|
||
|
||
def sess_of(sid):
|
||
with db.get_connection() as conn:
|
||
r = conn.execute("SELECT * FROM sessions WHERE id = ?", (sid,)).fetchone()
|
||
return dict(r) if r else None
|
||
|
||
|
||
def sess_count():
|
||
with db.get_connection() as conn:
|
||
return conn.execute("SELECT COUNT(*) FROM sessions").fetchone()[0]
|
||
|
||
|
||
def chain_sig(sid):
|
||
"""(role, content, is_ignored) 序列"""
|
||
return [(m["role"], m["content"], m["is_ignored"])
|
||
for m in db.get_message_chain(sid)]
|
||
|
||
|
||
# ======================================================================
|
||
try:
|
||
# ---------------- T1/T2: 线性会话 ----------------
|
||
src, src_ids = mk_session("线性会话", [
|
||
("system", "sys"), ("user", "u1"), ("assistant", "a1"),
|
||
("user", "u2"), ("assistant", "a2")])
|
||
n_before = sess_count()
|
||
cpy = db.copy_session(src)
|
||
check("T1.1 返回新会话 dict 且 id 不同", cpy and cpy["id"] != src)
|
||
check("T1.2 会话数 +1", sess_count() == n_before + 1, f"{n_before}->{sess_count()}")
|
||
|
||
s_rows, c_rows = rows_of(src), rows_of(cpy["id"])
|
||
check("T1.3 消息数一致", len(s_rows) == len(c_rows) == 5, f"{len(s_rows)}/{len(c_rows)}")
|
||
check("T1.4 ID 全新无交集",
|
||
not (set(r["id"] for r in s_rows) & set(r["id"] for r in c_rows)))
|
||
check("T1.5 role+content 序列一致",
|
||
[(r["role"], r["content"]) for r in s_rows] ==
|
||
[(r["role"], r["content"]) for r in c_rows])
|
||
check("T1.6 reasoning/timeline/usage 原样复制",
|
||
all(s["reasoning"] == c["reasoning"] and s["timeline"] == c["timeline"]
|
||
and s["usage"] == c["usage"] for s, c in zip(s_rows, c_rows)))
|
||
|
||
# T2: parent 链重建
|
||
idmap = {s["id"]: c["id"] for s, c in zip(s_rows, c_rows)}
|
||
ok_chain = c_rows[0]["parent_id"] is None
|
||
for s, c in list(zip(s_rows, c_rows))[1:]:
|
||
ok_chain = ok_chain and c["parent_id"] == idmap.get(s["parent_id"])
|
||
check("T2.1 parent_id 全量重映射", ok_chain)
|
||
check("T2.2 副本链首 parent 为 None（未指回源 ID）", c_rows[0]["parent_id"] is None)
|
||
check("T2.3 副本链里不含任何源 ID",
|
||
not (set(idmap.keys()) & set(r["parent_id"] for r in c_rows if r["parent_id"])))
|
||
|
||
# ---------------- T3: 分叉会话整树复制 ----------------
|
||
sid3, ids3 = mk_session("分叉会话", [("system", "s"), ("user", "q")])
|
||
base = ids3[-1]
|
||
add_child(sid3, base, "assistant", "br1", int(time.time()) + 10)
|
||
add_child(sid3, base, "assistant", "br2", int(time.time()) + 11)
|
||
copy3 = db.copy_session(sid3)
|
||
c3 = rows_of(copy3["id"])
|
||
src_branches = len(db.get_branch_info(base))
|
||
c_base = [r for r in c3 if r["role"] == "user"][0]["id"]
|
||
check("T3.1 源有 2 个分支", src_branches == 2, str(src_branches))
|
||
check("T3.2 副本分支数一致", len(db.get_branch_info(c_base)) == 2,
|
||
str(len(db.get_branch_info(c_base))))
|
||
check("T3.3 副本消息总数一致", len(c3) == len(rows_of(sid3)) == 4)
|
||
check("T3.4 分支内容一致",
|
||
sorted(b["content"] for b in db.get_branch_info(c_base)) == ["br1", "br2"])
|
||
|
||
# ---------------- T4: 压缩标记 ----------------
|
||
sid4, ids4 = mk_session("压缩会话", [
|
||
("system", "s"), ("user", "u1"), ("assistant", "a1"),
|
||
("user", "u2"), ("assistant", "a2")])
|
||
mark = db.insert_compaction_mark(sid4, "【摘要】前面聊了 u1/a1",
|
||
cut_before_id=ids4[2], first_retained_id=ids4[3])
|
||
copy4 = db.copy_session(sid4)
|
||
c4 = rows_of(copy4["id"])
|
||
s4 = rows_of(sid4)
|
||
c_marks = [r for r in c4 if r["role"] == "compaction"]
|
||
s_marks = [r for r in s4 if r["role"] == "compaction"]
|
||
check("T4.1 标记存在且已复制", len(s_marks) == 1 and len(c_marks) == 1)
|
||
check("T4.2 摘要内容一致", c_marks and c_marks[0]["content"] == s_marks[0]["content"])
|
||
m4 = {s["id"]: c["id"] for s, c in zip(s4, c4)}
|
||
check("T4.3 标记 parent 指向副本的 cut_before",
|
||
c_marks[0]["parent_id"] == m4.get(ids4[2]))
|
||
check("T4.4 first_retained 的 parent 指向副本标记",
|
||
[r for r in c4 if r["id"] == m4.get(ids4[3])][0]["parent_id"] == c_marks[0]["id"])
|
||
check("T4.5 标记 id 带 comp_ 前缀", c_marks[0]["id"].startswith("comp_"))
|
||
retained = [r for r in c4 if r["created_at"] is not None
|
||
and r["id"] != c_marks[0]["id"]]
|
||
check("T4.6 标记时间戳晚于全部其它行",
|
||
all(c_marks[0]["created_at"] >= r["created_at"] for r in retained))
|
||
check("T4.7 标记 is_ignored 保持 1", c_marks[0]["is_ignored"] == 1)
|
||
# created_at 相对次序必须与源完全一致（兄弟排序依赖它）
|
||
s_order = [i for i, _ in sorted(enumerate([r["created_at"] for r in s4]),
|
||
key=lambda x: x[1])]
|
||
c_order = [i for i, _ in sorted(enumerate([r["created_at"] for r in c4]),
|
||
key=lambda x: x[1])]
|
||
check("T4.8 created_at 相对次序与源一致", s_order == c_order, f"{s_order} vs {c_order}")
|
||
|
||
# ---------------- T5: 叶子重映射 ----------------
|
||
check("T5.1 副本叶子非空且不是源叶子",
|
||
cpy["current_leaf_msg_id"] and cpy["current_leaf_msg_id"] != sess_of(src)["current_leaf_msg_id"])
|
||
check("T5.2 副本叶子真实存在于副本",
|
||
cpy["current_leaf_msg_id"] in [r["id"] for r in c_rows])
|
||
check("T5.3 副本叶子对应源叶子的内容",
|
||
[r for r in c_rows if r["id"] == cpy["current_leaf_msg_id"]][0]["content"]
|
||
== s_rows[-1]["content"])
|
||
|
||
# ---------------- T6: 源会话零改动 ----------------
|
||
src_snapshot = (sess_of(src), [tuple(sorted(r.items())) for r in rows_of(src)])
|
||
_ = db.copy_session(src)
|
||
src_after = (sess_of(src), [tuple(sorted(r.items())) for r in rows_of(src)])
|
||
check("T6.1 复制两次后源 sessions 行不变", src_snapshot[0] == src_after[0])
|
||
check("T6.2 复制两次后源 messages 全列不变", src_snapshot[1] == src_after[1])
|
||
|
||
# ---------------- T7: 双向删除隔离 ----------------
|
||
sid7, _ = mk_session("隔离A", [("system", "s"), ("user", "x")])
|
||
c7 = db.copy_session(sid7)
|
||
db.delete_session(c7["id"])
|
||
check("T7.1 删副本后源仍在且消息完整",
|
||
sess_of(sid7) is not None and len(rows_of(sid7)) == 2)
|
||
# 方向二另起一对，避免与方向一互相干扰
|
||
sid7b, _ = mk_session("隔离B", [("system", "s"), ("user", "y")])
|
||
c7b = db.copy_session(sid7b)
|
||
db.delete_session(sid7b)
|
||
check("T7.2 删源后副本仍在且消息完整",
|
||
sess_of(c7b["id"]) is not None and len(rows_of(c7b["id"])) == 2,
|
||
f'cpy sess={sess_of(c7b["id"]) is not None} rows={len(rows_of(c7b["id"]))}')
|
||
|
||
# ---------------- T8: 图片附件物理复制 ----------------
|
||
png_src = os.path.join(_ATT, "img_src_test.png")
|
||
with open(png_src, "wb") as f:
|
||
f.write(b"\x89PNG\r\n\x1a\n" + b"FAKEIMAGEDATA" * 8)
|
||
img_meta = json.dumps({"user_text": "看看这张图",
|
||
"attachments": [{"type": "image", "size_kb": 0.2,
|
||
"local_path": "data/attachments/img_src_test.png"}]},
|
||
ensure_ascii=False)
|
||
sid8, _ = mk_session("图片会话", [("system", "s"), ("user", "带图")],
|
||
att_meta={1: img_meta})
|
||
c8 = db.copy_session(sid8)
|
||
m8 = [r for r in rows_of(c8["id"]) if r["attachment_metadata"]][0]
|
||
meta8 = json.loads(m8["attachment_metadata"])
|
||
new_rel = meta8["attachments"][0]["local_path"]
|
||
new_abs = os.path.join(_TMP, new_rel)
|
||
check("T8.1 local_path 已改写为新文件",
|
||
new_rel != "data/attachments/img_src_test.png", new_rel)
|
||
check("T8.2 新旧文件同时存在",
|
||
os.path.isfile(png_src) and os.path.isfile(new_abs))
|
||
check("T8.3 新文件内容与源一致",
|
||
open(new_abs, "rb").read() == open(png_src, "rb").read())
|
||
check("T8.4 user_text 等其它字段保留", meta8["user_text"] == "看看这张图")
|
||
check("T8.5 源 metadata 未被改动",
|
||
json.loads([r for r in rows_of(sid8) if r["attachment_metadata"]][0]
|
||
["attachment_metadata"])["attachments"][0]["local_path"]
|
||
== "data/attachments/img_src_test.png")
|
||
|
||
# 缺文件容错：metadata 保留原路径，不抛异常
|
||
bad_meta = json.dumps({"user_text": "x", "attachments": [
|
||
{"type": "image", "local_path": "data/attachments/does_not_exist.png"}]},
|
||
ensure_ascii=False)
|
||
sid8b, _ = mk_session("缺文件", [("system", "s"), ("user", "y")], att_meta={1: bad_meta})
|
||
c8b = db.copy_session(sid8b)
|
||
bad_after = [r for r in rows_of(c8b["id"]) if r["attachment_metadata"]][0]
|
||
check("T8.6 源文件缺失时保留原路径且不失败",
|
||
json.loads(bad_after["attachment_metadata"])["attachments"][0]["local_path"]
|
||
== "data/attachments/does_not_exist.png")
|
||
|
||
# ---------------- T9: 文本附件原样 ----------------
|
||
txt_meta = json.dumps({"user_text": "", "attachments": [
|
||
{"type": "text", "size_kb": 1.0, "lines": 3, "content": "aaa\nbbb\nccc"}]},
|
||
ensure_ascii=False)
|
||
sid9, _ = mk_session("文本附件", [("system", "s"), ("user", "t")], att_meta={1: txt_meta})
|
||
n_files_before = len(os.listdir(_ATT))
|
||
c9 = db.copy_session(sid9)
|
||
t9 = [r for r in rows_of(c9["id"]) if r["attachment_metadata"]][0]
|
||
check("T9.1 文本附件 metadata 逐字节相同",
|
||
t9["attachment_metadata"] == txt_meta, t9["attachment_metadata"][:60])
|
||
check("T9.2 文本附件不产生新文件", len(os.listdir(_ATT)) == n_files_before)
|
||
|
||
# ---------------- T10: 字段语义 ----------------
|
||
sid10, _ = mk_session("语义检查", [("system", "s"), ("user", "m")],
|
||
mode="worker", starred=1)
|
||
min_order_before = db.get_min_sort_order()
|
||
c10 = db.copy_session(sid10)
|
||
check("T10.1 标题 = 原名 + ' (副本)'", c10["title"] == "语义检查 (副本)", c10["title"])
|
||
check("T10.2 副本不带星标", c10["is_starred"] == 0, str(c10["is_starred"]))
|
||
check("T10.3 mode 跟随源", c10["mode"] == "worker", str(c10["mode"]))
|
||
check("T10.4 has_messages 跟随源", c10["has_messages"] == sess_of(sid10)["has_messages"])
|
||
check("T10.5 sort_order 置顶", c10["sort_order"] == min_order_before,
|
||
f'{c10["sort_order"]} vs {min_order_before}')
|
||
check("T10.6 源仍是星标且未被改标题",
|
||
sess_of(sid10)["is_starred"] == 1 and sess_of(sid10)["title"] == "语义检查")
|
||
check("T10.7 created_at/updated_at = 当前时间",
|
||
abs(c10["updated_at"] - int(time.time())) <= 5)
|
||
|
||
# 无消息的会话（只有 system 行）也能复制
|
||
sid10b, ids10b = mk_session("孤儿", [("system", "only")], has_messages=0)
|
||
c10b = db.copy_session(sid10b)
|
||
check("T10.8 仅 system 行的会话可复制",
|
||
c10b and len(rows_of(c10b["id"])) == 1 and c10b["current_leaf_msg_id"] is not None)
|
||
|
||
# ---------------- T12: 副本再复制 → (副本 2) ----------------
|
||
c12 = db.copy_session(c10["id"])
|
||
check("T12.1 再复制标题递增为 (副本 2)", c12["title"] == "语义检查 (副本 2)", c12["title"])
|
||
c12b = db.copy_session(c12["id"])
|
||
check("T12.2 第三次复制为 (副本 3)", c12b["title"] == "语义检查 (副本 3)", c12b["title"])
|
||
c12c = db.copy_session(sid10, new_title="自定义标题")
|
||
check("T12.3 可显式指定标题", c12c["title"] == "自定义标题", c12c["title"])
|
||
|
||
# ---------------- T11: 不存在的会话 ----------------
|
||
n11 = sess_count()
|
||
check("T11.1 不存在 → None", db.copy_session("sess_not_exist_xxx") is None)
|
||
check("T11.2 无残留（会话数不变）", sess_count() == n11)
|
||
|
||
# ---------------- T13: 外键 / 孤儿 ----------------
|
||
bad_fk = 0
|
||
orphans = 0
|
||
with db.get_connection() as conn:
|
||
bad_fk = len(conn.execute("PRAGMA foreign_key_check").fetchall())
|
||
for r in rows_of(cpy["id"]):
|
||
if r["parent_id"] and r["parent_id"] not in [x["id"] for x in c_rows]:
|
||
orphans += 1
|
||
check("T13.1 外键无违规", bad_fk == 0, str(bad_fk))
|
||
check("T13.2 无孤儿 parent_id", orphans == 0, str(orphans))
|
||
|
||
# ---------------- T14: 链签名一致（端到端） ----------------
|
||
check("T14.1 线性会话链签名一致", chain_sig(src) == chain_sig(cpy["id"]))
|
||
check("T14.2 分叉会话链签名一致", chain_sig(sid3) == chain_sig(copy3["id"]))
|
||
check("T14.3 压缩会话链签名一致（含标记行）", chain_sig(sid4) == chain_sig(copy4["id"]))
|
||
|
||
finally:
|
||
shutil.rmtree(_TMP, ignore_errors=True)
|
||
|
||
failed = [n for n, ok in RESULTS if not ok]
|
||
print(f"\n===== {len(RESULTS) - len(failed)}/{len(RESULTS)} PASS =====", flush=True)
|
||
print("ALL PASS" if not failed else f"FAILED: {failed}", flush=True)
|
||
sys.exit(0 if not failed else 1)
|