BossAi/test_chat_history.py

"""对话历史落库测试:验证普通/保密模式落库、加密、历史查询、阅后即焚不落库、密码验证。

运行: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()