152 lines
4.9 KiB
Python
152 lines
4.9 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""引擎 B 数据源模式单测：parallel / db / json 三模式调度 + DB 故障隔离。
|
||
|
||
验证点：
|
||
1. json 模式：每轮只跑 conversations.json（poll_once），不碰 DB；
|
||
2. db 模式（有 db_source）：只跑 DB 直读，不碰 JSON；
|
||
3. db 模式（无 db_source）：自动退化为 JSON，避免静默无检测；
|
||
4. parallel 模式：两路每轮都跑，互为兜底；
|
||
5. DB 故障隔离：get_new_messages 抛异常 → db_active=False，
|
||
JSON 路径下一轮仍独立检测投递；DB 恢复后 db_active 回 True。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
from engine_b import DataEngine # noqa: E402
|
||
|
||
FAILED: list[str] = []
|
||
|
||
|
||
def check(name: str, cond: bool, detail: str = "") -> None:
|
||
tag = "PASS" if cond else "FAIL"
|
||
print(f"[{tag}] {name}" + (f" ({detail})" if detail else ""))
|
||
if not cond:
|
||
FAILED.append(name)
|
||
|
||
|
||
class CountingEngine(DataEngine):
|
||
"""重写两路检测方法为计数器，验证 _run 的调度分支。"""
|
||
|
||
def __init__(self, *args, **kwargs):
|
||
super().__init__(*args, **kwargs)
|
||
self.poll_once_calls = 0
|
||
self.db_once_calls = 0
|
||
self.poll_lock = threading.Lock()
|
||
|
||
def poll_once(self):
|
||
with self.poll_lock:
|
||
self.poll_once_calls += 1
|
||
|
||
def _poll_db_once(self):
|
||
with self.poll_lock:
|
||
self.db_once_calls += 1
|
||
|
||
|
||
class StubDB:
|
||
def __init__(self, rows=None, fail=False):
|
||
self.rows = rows or []
|
||
self.fail = fail
|
||
self.calls = 0
|
||
|
||
def get_new_messages(self, since_ts):
|
||
self.calls += 1
|
||
if self.fail:
|
||
raise RuntimeError("DB 解密失败")
|
||
return self.rows
|
||
|
||
|
||
def run_rounds(engine: DataEngine, seconds: float = 1.5) -> None:
|
||
engine.start()
|
||
time.sleep(seconds)
|
||
engine.stop()
|
||
|
||
|
||
# ---------- 1/2/4. 三模式调度 ----------
|
||
mode_cases = [
|
||
("json", {}, {"poll": 2, "db": 0}),
|
||
("db", {"db_source": StubDB()}, {"poll": 0, "db": 2}),
|
||
("parallel", {"db_source": StubDB()}, {"poll": 2, "db": 2}),
|
||
]
|
||
for mode, kw, expect in mode_cases:
|
||
eng = CountingEngine(bot=object(), conversations_path="/nonexistent.json",
|
||
poll_interval=0.4, **kw)
|
||
eng.data_source_mode = mode
|
||
run_rounds(eng, seconds=1.3)
|
||
check(f"[{mode}] 每轮调用 poll_once 次数符合预期",
|
||
eng.poll_once_calls >= expect["poll"],
|
||
f"got {eng.poll_once_calls}")
|
||
check(f"[{mode}] 每轮调用 _poll_db_once 次数符合预期",
|
||
eng.db_once_calls >= expect["db"],
|
||
f"got {eng.db_once_calls}")
|
||
|
||
# ---------- 3. db 模式无数据源 → 退化为 JSON ----------
|
||
eng = CountingEngine(bot=object(), conversations_path="/nonexistent.json",
|
||
poll_interval=0.4, db_source=None)
|
||
eng.data_source_mode = "db"
|
||
run_rounds(eng, seconds=1.3)
|
||
check("[db-无数据源] 退化为 JSON 检测", eng.poll_once_calls >= 2,
|
||
f"poll_once={eng.poll_once_calls}")
|
||
|
||
# ---------- 5. DB 故障隔离与恢复 ----------
|
||
now = time.time()
|
||
fp = "f" * 40
|
||
conversations = {
|
||
fp: {
|
||
"display_name": "隔离测试客户",
|
||
"history": [{"role": "user", "content": "DB 挂了不影响我", "ts": now - 3}],
|
||
"last_lines": ["DB 挂了不影响我"],
|
||
}
|
||
}
|
||
tmp_json = os.path.join(tempfile.gettempdir(), "engine_b_mode_test.json")
|
||
with open(tmp_json, "w", encoding="utf-8") as handle:
|
||
json.dump(conversations, handle, ensure_ascii=False)
|
||
|
||
class StubBot:
|
||
def __init__(self):
|
||
self.enqueued = []
|
||
def has_active_pending(self, fp_hex):
|
||
return False
|
||
def enqueue_detected(self, **kwargs):
|
||
self.enqueued.append(kwargs)
|
||
return True, "ok"
|
||
|
||
bot = StubBot()
|
||
db = StubDB(fail=True)
|
||
eng = DataEngine(bot=bot, conversations_path=tmp_json,
|
||
poll_interval=0.4, db_source=db, data_source_mode="parallel")
|
||
eng._poll_db_once()
|
||
check("DB 异常 → db_active=False", eng.db_active is False,
|
||
f"db_active={eng.db_active}")
|
||
check("DB 异常 → last_error 记录", "DB 直读失败" in eng.last_error,
|
||
eng.last_error)
|
||
eng.poll_once()
|
||
check("DB 异常后 JSON 路径仍独立投递", bot.enqueued and eng.enqueued_count == 1,
|
||
f"enqueued={len(bot.enqueued)}")
|
||
check("DB 异常不传播到 _run（线程安全）",
|
||
True) # 上面的直接调用已证明 _poll_db_once 自吞异常
|
||
|
||
# DB 恢复
|
||
db.fail = False
|
||
eng._poll_db_once()
|
||
check("DB 恢复 → db_active 回 True", eng.db_active is True,
|
||
f"db_active={eng.db_active}")
|
||
|
||
# 非法模式回落
|
||
eng2 = DataEngine(bot=bot, conversations_path=tmp_json,
|
||
poll_interval=0.5, data_source_mode="garbage")
|
||
check("非法模式回落 parallel", eng2.data_source_mode == "parallel",
|
||
eng2.data_source_mode)
|
||
|
||
print()
|
||
if FAILED:
|
||
print(f"RESULT: FAIL ({len(FAILED)}) -> {FAILED}")
|
||
sys.exit(1)
|
||
print("RESULT: ALL PASS")
|