chongming-core/test/test_call_stats.py

"""
调用统计 + 审计日志（后台任务）测试（依赖隔离：桩掉 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.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]

# 审计采样率固定为全量：生产 config 的 access_sample_rate=0.1 下，成功请求的审计行
# 本就是按概率写的，断言"几条请求 = 几条审计行"会随机失败。本测试验证的是审计链路
# 是否工作，与采样策略无关，故固定为 1.0。
_cs_mod._AUDIT_SAMPLE_RATE = 1.0

# 加载真实 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,
                       instance_id=None):
        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,
                          instance_id=None):
        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("全部通过")