334 lines
13 KiB
Python
334 lines
13 KiB
Python
"""
|
||
调用统计 + 审计日志（后台任务）测试（依赖隔离：桩掉 chongming 包链、auth、
|
||
chongming_cache 与 worker_manager，加载真实 call_stats 与 route_generator）
|
||
|
||
用 httpx.ASGITransport 验证:
|
||
1. CallStats.record/summary 聚合正确性（计数 / 平均 / 最大延迟 / recent 上限）
|
||
2. spawn/shutdown：后台任务被跟踪、shutdown 后全部取消
|
||
3. 动态端点请求后：响应与改动前一致（字节透传）+ 统计计数正确 + 审计日志出现 AUDIT 行
|
||
4. 错误路径（无可用 worker）-> HTTP 503 同样被记录
|
||
5. 后台任务失败兜底：record_call 内部异常不冒泡，端点响应不受影响
|
||
|
||
运行: uv run python test_call_stats.py
|
||
"""
|
||
import asyncio
|
||
import importlib.util
|
||
import logging
|
||
import sys
|
||
import types
|
||
from pathlib import Path
|
||
|
||
import httpx
|
||
import orjson
|
||
from fastapi import FastAPI
|
||
|
||
REPO_ROOT = Path(__file__).resolve().parent
|
||
sys.path.insert(0, str(REPO_ROOT / "src"))
|
||
sys.path.insert(0, str(REPO_ROOT / "utils" / "worker" / "src"))
|
||
|
||
# ---- 依赖隔离：桩掉 chongming 包链 / auth / chongming_cache / worker_manager ----
|
||
_pkg = types.ModuleType("chongming")
|
||
_pkg.__path__ = []
|
||
_app = types.ModuleType("chongming.app")
|
||
_app.__path__ = []
|
||
_core = types.ModuleType("chongming.app.core")
|
||
_core.__path__ = []
|
||
_logger_stub = types.ModuleType("chongming.app.core.logger")
|
||
_logger_stub.get_logger = lambda name: logging.getLogger(name)
|
||
_auth_stub = types.ModuleType("chongming.app.core.auth")
|
||
_auth_stub.store_user_to_request = lambda x: x # 仅用于 Depends，桩掉鉴权
|
||
_cache_stub = types.ModuleType("chongming_cache")
|
||
_cache_stub.Cache = object
|
||
# worker_manager 桩：route_generator 仅用其类型注解与异常类
|
||
_wm_stub = types.ModuleType("chongming.app.core.worker_manager")
|
||
_wm_stub.WorkerManager = object
|
||
_wm_stub.StreamIdleTimeoutError = type("StreamIdleTimeoutError", (RuntimeError,), {})
|
||
for mod in (_pkg, _app, _core, _logger_stub, _auth_stub, _cache_stub, _wm_stub):
|
||
sys.modules[mod.__name__] = mod
|
||
|
||
# 加载真实 call_stats（route_generator 依赖其 spawn/record_call）
|
||
_cs_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "call_stats.py"
|
||
_cs_spec = importlib.util.spec_from_file_location("chongming.app.core.call_stats", _cs_path)
|
||
_cs_mod = importlib.util.module_from_spec(_cs_spec)
|
||
sys.modules["chongming.app.core.call_stats"] = _cs_mod
|
||
_cs_spec.loader.exec_module(_cs_mod) # type: ignore[union-attr]
|
||
|
||
# 加载真实 route_generator
|
||
_rg_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "route_generator.py"
|
||
_rg_spec = importlib.util.spec_from_file_location("chongming.app.core.route_generator", _rg_path)
|
||
_rg_mod = importlib.util.module_from_spec(_rg_spec)
|
||
sys.modules["chongming.app.core.route_generator"] = _rg_mod
|
||
_rg_spec.loader.exec_module(_rg_mod) # type: ignore[union-attr]
|
||
|
||
register_route_group = _rg_mod.register_route_group
|
||
record_call = _cs_mod.record_call
|
||
spawn = _cs_mod.spawn
|
||
|
||
# ---- 审计日志捕获（桩 logger 直接挂到标准库 audit logger） ----
|
||
_audit_records = []
|
||
|
||
|
||
class _AuditHandler(logging.Handler):
|
||
def emit(self, record):
|
||
_audit_records.append(record.getMessage())
|
||
|
||
|
||
_audit_logger = logging.getLogger("audit")
|
||
_audit_logger.setLevel(logging.INFO)
|
||
_audit_logger.addHandler(_AuditHandler())
|
||
_audit_logger.propagate = False
|
||
|
||
failures = []
|
||
|
||
|
||
def check(name, cond, detail=""):
|
||
status = "PASS" if cond else "FAIL"
|
||
print(f"[{status}] {name}" + (f" — {detail}" if detail else ""))
|
||
if not cond:
|
||
failures.append(name)
|
||
|
||
|
||
class FakeWorkerManager:
|
||
"""模拟 WorkerManager：call_raw 返回原始字节；error="no_worker" 时抛 RuntimeError"""
|
||
|
||
def __init__(self, reply=b'{"ok": true}', error=None):
|
||
self.reply = reply
|
||
self.error = error
|
||
|
||
async def call_raw(self, worker_type, subject, payload, timeout=3.0, retries=1, cacheable=False):
|
||
if self.error == "no_worker":
|
||
raise RuntimeError(f"类型 '{worker_type}' 没有可用 Worker。")
|
||
return self.reply
|
||
|
||
async def call_stream(self, worker_type, subject, payload, timeout=300.0, idle_timeout=10.0):
|
||
if self.error == "no_worker":
|
||
raise RuntimeError(f"类型 '{worker_type}' 没有可用 Worker。")
|
||
|
||
async def _gen():
|
||
yield orjson.dumps({"chunk": 1})
|
||
yield orjson.dumps({"chunk": 2})
|
||
|
||
return _gen() # 与真实实现一致：async 函数返回异步生成器
|
||
|
||
|
||
ROUTES = [
|
||
# body 端点（validate=true）
|
||
{"subject": "demo.echo", "path": "/echo", "method": "POST",
|
||
"params": ["msg: str"], "timeout": 3.0, "auth_required": False},
|
||
# raw body 端点（validate=false）
|
||
{"subject": "demo.raw", "path": "/raw/echo", "method": "POST",
|
||
"params": ["msg: str"], "validate": False, "timeout": 3.0, "auth_required": False},
|
||
# query 带参端点
|
||
{"subject": "demo.add", "path": "/add", "method": "GET",
|
||
"params": ["a: int", "b: int"], "timeout": 3.0, "auth_required": False},
|
||
# query 无参端点（_simple_endpoint 分支）
|
||
{"subject": "demo.ping", "path": "/ping", "method": "GET",
|
||
"timeout": 3.0, "auth_required": False},
|
||
# 流式端点
|
||
{"subject": "demo.stream", "path": "/stream", "method": "GET",
|
||
"stream": True, "timeout": 3.0, "auth_required": False},
|
||
]
|
||
|
||
|
||
def build_app(fwm):
|
||
app = FastAPI()
|
||
register_route_group(app, fwm, "demo", ROUTES, prefix="/api/v1")
|
||
return app
|
||
|
||
|
||
async def _drain_background():
|
||
"""等待并清理所有后台统计/审计任务（每个异步测试末尾调用）"""
|
||
tasks = list(_cs_mod._tasks)
|
||
if tasks:
|
||
await asyncio.gather(*tasks, return_exceptions=True)
|
||
await _cs_mod.shutdown()
|
||
|
||
|
||
def _reset_stats():
|
||
"""隔离：每个端点测试前重置全局统计实例"""
|
||
_cs_mod._stats = _cs_mod.CallStats()
|
||
|
||
|
||
# ---- 1. CallStats 单元 ----
|
||
|
||
def test_stats_aggregation():
|
||
cs = _cs_mod.CallStats(max_recent=10)
|
||
cs.record({"worker_type": "demo", "subject": "demo.echo", "status": 200, "duration_ms": 10.0})
|
||
cs.record({"worker_type": "demo", "subject": "demo.echo", "status": 200, "duration_ms": 30.0})
|
||
cs.record({"worker_type": "demo", "subject": "demo.echo", "status": 503, "duration_ms": 5.0})
|
||
cs.record({"worker_type": "other", "subject": "other.x", "status": 200, "duration_ms": 2.0})
|
||
s = cs.summary()
|
||
check("total_calls", s["total_calls"] == 4, str(s["total_calls"]))
|
||
by_key = {(c["worker_type"], c["subject"], c["status"]): c["count"] for c in s["calls"]}
|
||
check(
|
||
"per-key 计数",
|
||
by_key == {
|
||
("demo", "demo.echo", 200): 2,
|
||
("demo", "demo.echo", 503): 1,
|
||
("other", "other.x", 200): 1,
|
||
},
|
||
str(by_key),
|
||
)
|
||
lat = {l["worker_type"]: l for l in s["latency"]}
|
||
check("延迟聚合 count", lat["demo"]["count"] == 3, str(lat["demo"]))
|
||
check("延迟聚合 avg", lat["demo"]["avg_ms"] == 15.0, str(lat["demo"]["avg_ms"]))
|
||
check("延迟聚合 max", lat["demo"]["max_ms"] == 30.0, str(lat["demo"]["max_ms"]))
|
||
check("recent 保留", len(s["recent"]) == 4, str(len(s["recent"])))
|
||
|
||
|
||
def test_recent_bounded():
|
||
cs = _cs_mod.CallStats(max_recent=5)
|
||
for _ in range(20):
|
||
cs.record({"worker_type": "demo", "subject": "demo.x", "status": 200, "duration_ms": 1.0})
|
||
s = cs.summary()
|
||
check("recent 环形缓冲上限", len(s["recent"]) == 5, str(len(s["recent"])))
|
||
check("计数不受缓冲影响", s["total_calls"] == 20, str(s["total_calls"]))
|
||
|
||
|
||
# ---- 2. spawn / shutdown ----
|
||
|
||
def test_spawn_shutdown():
|
||
async def run():
|
||
async def never():
|
||
await asyncio.sleep(3600)
|
||
|
||
t = spawn(never())
|
||
check("任务被跟踪", t in _cs_mod._tasks, str(list(_cs_mod._tasks)))
|
||
await _cs_mod.shutdown()
|
||
check("shutdown 后任务取消", t.cancelled(), str(t))
|
||
check("任务集合清空", len(_cs_mod._tasks) == 0, str(len(_cs_mod._tasks)))
|
||
|
||
asyncio.run(run())
|
||
|
||
|
||
# ---- 3. 端点集成：统计 + 审计 ----
|
||
|
||
def test_endpoint_records_stats_and_audit():
|
||
async def run():
|
||
_audit_records.clear()
|
||
_reset_stats()
|
||
fwm = FakeWorkerManager(reply=b'{"ok": true}')
|
||
app = build_app(fwm)
|
||
transport = httpx.ASGITransport(app=app)
|
||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as ac:
|
||
r1 = await ac.post("/api/v1/echo", json={"msg": "hi"})
|
||
r2 = await ac.post("/api/v1/raw/echo", json={"msg": "hi"})
|
||
r3 = await ac.get("/api/v1/add?a=1&b=2")
|
||
r4 = await ac.get("/api/v1/ping")
|
||
r5 = await ac.get("/api/v1/stream")
|
||
await _drain_background()
|
||
|
||
# 响应与改动前一致：worker 原始字节零再序列化
|
||
check("POST 响应字节透传", r1.status_code == 200 and r1.content == b'{"ok": true}', repr(r1.content))
|
||
check("raw POST 响应字节透传", r2.status_code == 200 and r2.content == b'{"ok": true}', repr(r2.content))
|
||
check("GET 响应字节透传", r3.status_code == 200 and r3.content == b'{"ok": true}', repr(r3.content))
|
||
check("无参 GET 响应字节透传", r4.status_code == 200 and r4.content == b'{"ok": true}', repr(r4.content))
|
||
check("流式响应 SSE 帧", r5.status_code == 200 and b"data:" in r5.content and b'"chunk":1' in r5.content, repr(r5.content[:80]))
|
||
|
||
s = _cs_mod.summary()
|
||
check("total_calls == 5", s["total_calls"] == 5, str(s["total_calls"]))
|
||
subjects = {c["subject"] for c in s["calls"]}
|
||
check(
|
||
"5 种调用均记录",
|
||
subjects == {"demo.echo", "demo.raw", "demo.add", "demo.ping", "demo.stream"},
|
||
str(subjects),
|
||
)
|
||
check("全部状态 200", all(c["status"] == 200 for c in s["calls"]), str(s["calls"]))
|
||
|
||
check("审计行数 == 5", len(_audit_records) == 5, str(len(_audit_records)))
|
||
echo_audit = next((a for a in _audit_records if "subject=demo.echo" in a), "")
|
||
check(
|
||
"审计行格式",
|
||
all(k in echo_audit for k in (
|
||
"AUDIT", "ts=", "worker=demo", "method=POST",
|
||
"path=/api/v1/echo", "status=200", "req_bytes=", "resp_bytes=",
|
||
)),
|
||
echo_audit,
|
||
)
|
||
stream_audit = next((a for a in _audit_records if "subject=demo.stream" in a), "")
|
||
check("流式审计已记录", "status=200" in stream_audit and "path=/api/v1/stream" in stream_audit, stream_audit)
|
||
|
||
asyncio.run(run())
|
||
|
||
|
||
def test_error_503_recorded():
|
||
async def run():
|
||
_audit_records.clear()
|
||
_reset_stats()
|
||
fwm = FakeWorkerManager(error="no_worker")
|
||
app = build_app(fwm)
|
||
transport = httpx.ASGITransport(app=app)
|
||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as ac:
|
||
r = await ac.get("/api/v1/ping")
|
||
await _drain_background()
|
||
|
||
check("503 响应", r.status_code == 503, str(r.status_code))
|
||
s = _cs_mod.summary()
|
||
check(
|
||
"503 计数 1",
|
||
any(c["subject"] == "demo.ping" and c["status"] == 503 and c["count"] == 1 for c in s["calls"]),
|
||
str(s["calls"]),
|
||
)
|
||
check(
|
||
"503 审计行",
|
||
any("subject=demo.ping" in a and "status=503" in a for a in _audit_records),
|
||
str(_audit_records),
|
||
)
|
||
|
||
asyncio.run(run())
|
||
|
||
|
||
# ---- 4. 后台任务失败兜底 ----
|
||
|
||
def test_background_failure_isolated():
|
||
async def run():
|
||
_audit_records.clear()
|
||
_reset_stats()
|
||
original = _cs_mod._stats.record
|
||
|
||
def boom(entry):
|
||
raise ValueError("stats broken")
|
||
|
||
_cs_mod._stats.record = boom
|
||
try:
|
||
# 后台执行体自身不冒泡（统计故障仅记 warning，审计仍写入）
|
||
await record_call({
|
||
"ts": 0.0, "worker_type": "demo", "subject": "demo.x", "status": 200,
|
||
"duration_ms": 1.0, "method": "GET", "path": "/x",
|
||
"user": None, "client": None, "req_bytes": 0, "resp_bytes": 0,
|
||
})
|
||
check("record_call 异常被兜底（不冒泡）", True)
|
||
check("审计仍写入", any("subject=demo.x" in a for a in _audit_records), str(_audit_records))
|
||
finally:
|
||
_cs_mod._stats.record = original
|
||
|
||
# 统计故障下端点响应不受影响
|
||
fwm = FakeWorkerManager(reply=b'{"ok": true}')
|
||
app = build_app(fwm)
|
||
_cs_mod._stats.record = boom
|
||
try:
|
||
transport = httpx.ASGITransport(app=app)
|
||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as ac:
|
||
r = await ac.get("/api/v1/ping")
|
||
await _drain_background()
|
||
check("统计故障下响应仍 200", r.status_code == 200 and r.content == b'{"ok": true}', repr(r.content))
|
||
finally:
|
||
_cs_mod._stats.record = original
|
||
|
||
asyncio.run(run())
|
||
|
||
|
||
# ---- 执行 ----
|
||
test_stats_aggregation()
|
||
test_recent_bounded()
|
||
test_endpoint_records_stats_and_audit()
|
||
test_error_503_recorded()
|
||
test_background_failure_isolated()
|
||
test_spawn_shutdown() # 最后执行：shutdown 会取消所有未完成任务
|
||
|
||
print()
|
||
if failures:
|
||
print(f"共 {len(failures)} 项失败: {failures}")
|
||
sys.exit(1)
|
||
print("全部通过")
|