"""对话历史落库测试:验证普通/保密模式落库、加密、历史查询、阅后即焚不落库、密码验证。 运行:python3 test_chat_history.py 前置:PostgreSQL 已启动,db_init.py 已执行。 """ import sys import psycopg2 from bossai import ChatStore, CryptoUtil, SessionManager def get_conn(): return psycopg2.connect(host="localhost", port=5432, dbname="mcp", user="postgres") def cleanup(): """清理测试数据。""" with get_conn() as conn: with conn.cursor() as cur: cur.execute("DELETE FROM chat_messages WHERE phone LIKE '139_test_%'") print("[cleanup] cleared test data") def test_crypto(): """测试加密解密往返。""" crypto = CryptoUtil() plaintext = "这是一条保密消息,包含敏感信息:手机号13800001111" encrypted = crypto.encrypt(plaintext) assert encrypted != plaintext, "加密后应与原文不同" decrypted = crypto.decrypt(encrypted) assert decrypted == plaintext, f"解密后应与原文一致: {decrypted}" print("[pass] 加密解密往返正确") def test_normal_mode_persist(): """测试普通模式落库(明文)。""" store = ChatStore() phone = "139_test_normal" sid = "test_session_normal_001" # 保存用户消息 msg_id = store.save_message( session_id=sid, phone=phone, mode="normal", role="user", content="员工旷工3天能开除吗?", scene="legal", ) assert msg_id is not None, "普通模式应返回消息 ID" # 保存 AI 回复(含风险卡) msg_id2 = store.save_message( session_id=sid, phone=phone, mode="normal", role="assistant", content="需要先确认考勤制度是否公示...", scene="legal", risk_level="amber", risk_conclusion="需补充考勤制度公示证据", risk_actions=["收集考勤记录", "确认制度公示"], risk_evidence=["考勤表"], risk_escalation=False, inference_location="cloud", model_name="moonshot-v1-8k", ) assert msg_id2 is not None # 验证落库 sessions = store.list_sessions(phone) assert len(sessions) == 1, f"应有1个会话,实际{len(sessions)}" assert sessions[0]["mode"] == "normal" assert sessions[0]["message_count"] == 2 # 验证消息内容(明文) messages = store.get_messages(sid, phone, decrypt=False) assert len(messages) == 2 assert messages[0]["role"] == "user" assert messages[0]["content"] == "员工旷工3天能开除吗?" assert messages[1]["role"] == "assistant" assert messages[1]["risk"]["level"] == "amber" assert messages[1]["risk"]["conclusion"] == "需补充考勤制度公示证据" assert messages[1]["meta"]["inference_location"] == "cloud" print("[pass] 普通模式落库正确(明文存储)") def test_confidential_mode_encrypt(): """测试保密模式落库(加密)。""" store = ChatStore() phone = "139_test_conf" sid = "test_session_conf_001" sensitive_content = "股东纠纷:我占股40%,大股东挪用资金200万" # 保存用户消息(保密模式) store.save_message( session_id=sid, phone=phone, mode="confidential", role="user", content=sensitive_content, scene="legal", ) # 保存 AI 回复 store.save_message( session_id=sid, phone=phone, mode="confidential", role="assistant", content="建议立即保全证据...", scene="legal", risk_level="red", risk_conclusion="高风险:可能涉及职务侵占", risk_actions=["保全财务凭证", "咨询律师"], risk_escalation=True, ) # 验证数据库中存储的是密文 with get_conn() as conn: with conn.cursor() as cur: cur.execute( "SELECT content FROM chat_messages WHERE session_id = %s AND role = 'user'", (sid,), ) row = cur.fetchone() assert row is not None stored_content = row[0] assert stored_content != sensitive_content, "保密模式数据库中应存储密文,非明文" print(f"[info] 明文: {sensitive_content[:20]}...") print(f"[info] 密文: {stored_content[:40]}...") # 验证不解密时返回密文 messages_raw = store.get_messages(sid, phone, decrypt=False) assert messages_raw[0]["content"] != sensitive_content, "未解密时应为密文" # 验证解密后返回明文 messages_decrypted = store.get_messages(sid, phone, decrypt=True) assert messages_decrypted[0]["content"] == sensitive_content, "解密后应为明文" assert messages_decrypted[1]["risk"]["conclusion"] == "高风险:可能涉及职务侵占" assert messages_decrypted[1]["risk"]["level"] == "red" assert messages_decrypted[1]["risk"]["escalation"] is True print("[pass] 保密模式加密存储正确") def test_ephemeral_not_persist(): """测试阅后即焚不落库。""" store = ChatStore() phone = "139_test_eph" sid = "test_session_eph_001" result = store.save_message( session_id=sid, phone=phone, mode="ephemeral", role="user", content="家庭隐私问题,阅后即焚", scene="emotion", ) assert result is None, "阅后即焚模式应返回 None(不落库)" # 验证数据库中没有该会话的数据 with get_conn() as conn: with conn.cursor() as cur: cur.execute( "SELECT COUNT(*) FROM chat_messages WHERE session_id = %s", (sid,) ) count = cur.fetchone()[0] assert count == 0, f"阅后即焚不应落库,实际有{count}条" print("[pass] 阅后即焚不落库正确") def test_list_sessions_filter(): """测试按模式过滤会话列表。""" store = ChatStore() phone = "139_test_filter" # 创建两种模式的会话 store.save_message(session_id="test_filter_n1", phone=phone, mode="normal", role="user", content="普通1") store.save_message(session_id="test_filter_c1", phone=phone, mode="confidential", role="user", content="保密1") # 全部 all_sessions = store.list_sessions(phone) assert len(all_sessions) == 2, f"应有2个会话,实际{len(all_sessions)}" # 只看普通 normal_sessions = store.list_sessions(phone, mode="normal") assert len(normal_sessions) == 1 assert normal_sessions[0]["mode"] == "normal" # 只看保密 conf_sessions = store.list_sessions(phone, mode="confidential") assert len(conf_sessions) == 1 assert conf_sessions[0]["mode"] == "confidential" print("[pass] 会话列表过滤正确") def test_delete_session(): """测试删除会话历史。""" store = ChatStore() phone = "139_test_delete" sid = "test_session_delete_001" store.save_message(session_id=sid, phone=phone, mode="normal", role="user", content="待删除") store.save_message(session_id=sid, phone=phone, mode="normal", role="assistant", content="回复") # 删除前确认存在 messages = store.get_messages(sid, phone) assert len(messages) == 2 # 删除 deleted = store.delete_session(sid, phone) assert deleted == 2, f"应删除2条,实际{deleted}" # 删除后确认不存在 messages_after = store.get_messages(sid, phone) assert len(messages_after) == 0 print("[pass] 删除会话历史正确") def test_cross_user_isolation(): """测试用户间数据隔离。""" store = ChatStore() phone_a = "139_test_iso_a" phone_b = "139_test_iso_b" sid_a = "test_iso_a_001" store.save_message(session_id=sid_a, phone=phone_a, mode="normal", role="user", content="用户A的消息") # 用户B不能看到用户A的会话 messages_b = store.get_messages(sid_a, phone_b) assert len(messages_b) == 0, "用户B不应能查看用户A的会话" print("[pass] 用户间数据隔离正确") def main(): cleanup() try: test_crypto() test_normal_mode_persist() test_confidential_mode_encrypt() test_ephemeral_not_persist() test_list_sessions_filter() test_delete_session() test_cross_user_isolation() print("\n===== 全部 7 项测试通过 =====") except AssertionError as e: print(f"\n[FAIL] {e}") sys.exit(1) finally: cleanup() if __name__ == "__main__": main()