194 lines
7.1 KiB
Python
194 lines
7.1 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""双引擎并存集成测试。
|
|
|
|
覆盖：发送互斥锁、B 引擎投递（enqueue_detected）、队列字段透传（load/persist
|
|
往返）、has_active_pending 去重判断、detected_by 合并、DataEngine 真实集成。
|
|
用法：python test_dual_engine.py
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
|
|
_ROOT = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.insert(0, _ROOT)
|
|
|
|
import send_lock
|
|
import engine_b
|
|
import wechat_bot
|
|
|
|
_PASS = 0
|
|
_FAIL = 0
|
|
|
|
|
|
def check(name, condition, detail=""):
|
|
global _PASS, _FAIL
|
|
if condition:
|
|
_PASS += 1
|
|
print(f" [PASS] {name}")
|
|
else:
|
|
_FAIL += 1
|
|
print(f" [FAIL] {name} {detail}")
|
|
|
|
|
|
def make_mini_bot():
|
|
"""构造最小 WeChatBot 对象：__new__ 绕过 1.3 万行 __init__，只补队列所需字段。"""
|
|
bot = wechat_bot.WeChatBot.__new__(wechat_bot.WeChatBot)
|
|
bot._pending_lock = threading.RLock()
|
|
bot._pending_reply_sessions = {}
|
|
bot._pending_reply_path = os.path.join(
|
|
tempfile.gettempdir(), f"test_pending_{os.getpid()}.json"
|
|
)
|
|
bot._cancelled_reply_sessions = set()
|
|
bot.identity_by_name = True
|
|
|
|
def fake_mark_reply_pending(fp, **kwargs):
|
|
"""模拟真实 _mark_reply_pending_unlocked 的 setdefault + display_name 语义。"""
|
|
key = fp.hex()
|
|
if key in bot._cancelled_reply_sessions:
|
|
return False
|
|
state = bot._pending_reply_sessions.setdefault(key, {
|
|
"batch_ready": False,
|
|
"confirmed_unread": False,
|
|
"requires_visual_proof": False,
|
|
"visual_rejection_count": 0,
|
|
"created_at": time.time(),
|
|
"stage": "queued",
|
|
"stage_started_at": time.time(),
|
|
"stage_history": [],
|
|
})
|
|
name = str(kwargs.get("display_name") or "").strip()
|
|
if name and not state.get("display_name"):
|
|
state["display_name"] = name
|
|
return True
|
|
|
|
bot._mark_reply_pending = fake_mark_reply_pending
|
|
return bot
|
|
|
|
|
|
def test_send_lock():
|
|
print("[1] 发送互斥锁")
|
|
assert send_lock.try_acquire("test_a", 0), "应能拿到锁"
|
|
held = send_lock.try_acquire("test_b", 0)
|
|
check("互斥：持锁期间他人抢不到", not held)
|
|
send_lock.release()
|
|
check("释放后他人可抢到", send_lock.try_acquire("test_c", 0))
|
|
send_lock.release()
|
|
# 上下文管理器
|
|
with send_lock.lock("test_d") as ok:
|
|
check("上下文管理器拿锁", ok)
|
|
check("持锁期间互斥", not send_lock.try_acquire("test_e", 0))
|
|
check("退出上下文已释放", send_lock.try_acquire("test_f", 0))
|
|
send_lock.release()
|
|
|
|
|
|
def test_enqueue_and_persist():
|
|
print("[2] B 引擎投递 + 队列字段透传")
|
|
bot = make_mini_bot()
|
|
fp_hex = "ab" * 40 # 40 字节会话指纹
|
|
ok, reason = bot.enqueue_detected(
|
|
fp_hex=fp_hex,
|
|
dedup_key=f"{fp_hex}:1234.5",
|
|
detected_by="engine_b",
|
|
chat_text="我想挂号",
|
|
display_name="测试客户",
|
|
last_lines=["我想挂号"],
|
|
)
|
|
check("首次投递成功", ok, reason)
|
|
state = bot._pending_reply_sessions.get(fp_hex)
|
|
check("条目已创建", state is not None)
|
|
if state:
|
|
check("dedup_key 已写入", state.get("dedup_key") == f"{fp_hex}:1234.5")
|
|
check("detected_by=engine_b", state.get("detected_by") == "engine_b")
|
|
check("chat_text 已写入", state.get("chat_text") == "我想挂号")
|
|
check("display_name 已写入", state.get("display_name") == "测试客户")
|
|
check("stage=queued", state.get("stage") == "queued")
|
|
|
|
# 重复投递同一会话：只补来源标记，不覆盖
|
|
ok, reason = bot.enqueue_detected(
|
|
fp_hex=fp_hex,
|
|
dedup_key=f"{fp_hex}:1234.5",
|
|
detected_by="engine_a",
|
|
chat_text="",
|
|
)
|
|
check("重复投递仍成功", ok, reason)
|
|
state = bot._pending_reply_sessions.get(fp_hex)
|
|
check("detected_by 合并为双引擎", state.get("detected_by") == "engine_b|engine_a", state.get("detected_by"))
|
|
check("chat_text 未被覆盖", state.get("chat_text") == "我想挂号")
|
|
|
|
# 持久化往返：新字段不丢
|
|
ok = bot._persist_pending_replies()
|
|
check("持久化成功", ok)
|
|
reloaded = wechat_bot.WeChatBot._load_pending_replies(bot)
|
|
restored = reloaded.get(fp_hex)
|
|
check("磁盘恢复条目", restored is not None)
|
|
if restored:
|
|
check("恢复后 dedup_key 保留", restored.get("dedup_key") == f"{fp_hex}:1234.5")
|
|
check("恢复后 detected_by 保留", restored.get("detected_by") == "engine_b|engine_a")
|
|
check("恢复后 detection_ts>0", float(restored.get("detection_ts") or 0) > 0)
|
|
check("恢复后 sending_owner 为空", restored.get("sending_owner") == "")
|
|
check("恢复后 send_lock_ts=0", float(restored.get("send_lock_ts") or 0) == 0)
|
|
|
|
# has_active_pending
|
|
check("活跃任务判断 True", bot.has_active_pending(fp_hex))
|
|
bot._pending_reply_sessions[fp_hex]["stage"] = "sent"
|
|
check("终态任务判断 False", not bot.has_active_pending(fp_hex))
|
|
bot._pending_reply_sessions[fp_hex]["stage"] = "queued"
|
|
|
|
# merge_detected_by
|
|
bot.merge_detected_by(fp_hex, "engine_a")
|
|
state = bot._pending_reply_sessions.get(fp_hex)
|
|
check("merge 幂等（不重复追加）", state.get("detected_by") == "engine_b|engine_a")
|
|
try:
|
|
os.unlink(bot._pending_reply_path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def test_engine_b_integration():
|
|
print("[3] DataEngine 与真实 bot 集成")
|
|
bot = make_mini_bot()
|
|
now = time.time()
|
|
conv_path = os.path.join(tempfile.gettempdir(), f"test_convs_{os.getpid()}.json")
|
|
with open(conv_path, "w", encoding="utf-8") as handle:
|
|
json.dump({
|
|
"cd" * 40: {
|
|
"display_name": "集成客户",
|
|
"history": [
|
|
{"role": "user", "content": "在吗", "ts": now - 30},
|
|
{"role": "assistant", "content": "在的", "ts": now - 29},
|
|
{"role": "user", "content": "帮我查下挂号的号源", "ts": now - 3},
|
|
],
|
|
"last_lines": ["帮我查下挂号的号源"],
|
|
}
|
|
}, handle, ensure_ascii=False)
|
|
engine = engine_b.DataEngine(bot=bot, conversations_path=conv_path, poll_interval=0.5)
|
|
engine.poll_once()
|
|
fp_hex = "cd" * 40
|
|
state = bot._pending_reply_sessions.get(fp_hex)
|
|
check("B 引擎投递到真实队列", state is not None)
|
|
if state:
|
|
check("dedup_key 含时间戳", state.get("dedup_key") == f"{fp_hex}:{now - 3}")
|
|
check("detected_by=engine_b", state.get("detected_by") == "engine_b")
|
|
check("chat_text 为增量文本", state.get("chat_text") == "帮我查下挂号的号源")
|
|
check("display_name 来自档案", state.get("display_name") == "集成客户")
|
|
# 再次轮询不重复投递
|
|
engine.poll_once()
|
|
check("重复轮询不重复投递", engine.enqueued_count == 1)
|
|
engine.stop()
|
|
try:
|
|
os.unlink(conv_path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_send_lock()
|
|
test_enqueue_and_persist()
|
|
test_engine_b_integration()
|
|
print(f"\n结果: {_PASS} 通过, {_FAIL} 失败")
|
|
sys.exit(1 if _FAIL else 0)
|