feat: Host 中间件租户路由（纯ASGI覆盖WS/未知404/停用503暂停页/health豁免）

Co-Authored-By: Claude Code <noreply@anthropic.com>
2026-10-12 20:30:30 +08:00
parent 1824bb1590
commit eaa5bf7e51
5 changed files with 338 additions and 50 deletions
@@ -3,9 +3,11 @@ import os
 import shutil
 import tempfile
 import warnings
+from contextlib import contextmanager
 from pathlib import Path

 import pytest
+from fastapi.testclient import TestClient
 from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
 from sqlalchemy.pool import NullPool

@@ -94,33 +96,87 @@ def build_jiancai_runtime():
    return _JIANCAI_RUNTIME


-@pytest.fixture
-def client(tmp_path):
-    """TestClient + 每测试独立 Registry(tmp_path 作 data_dir)。
+# Task 7: app_client 默认域名——既有测试不带 host 头,默认落 jiancai
+# (与 5b 之前 default_runtime 兜底的语义一致,既有断言零改动)
+JIANCAI_HOST = "chat2.agent.hnxinyouhe.com"
+AUTO_HOST = "chat.agent.hnxinyouhe.com"

-    - 预置 app.state.registry(测试注入缝): lifespan 见 registry 已挂则
-      沿用,不再加载仓库 seed_profiles/tenants——每测试拿到全新租户库,
-      与生产数据互不污染。
-    - 默认租户(jiancai)的 engine 由 lifespan 在 TestClient 的 portal 事件
-      循环内预热,后续请求也在同一循环处理(连接/循环绑定安全)。测试
-      线程需要直接造数时不得用该 engine——asyncio.run 另起循环,跨循环
-      复用常驻连接会死锁;直连断言走 async_db,或经真实路由(WS/REST)。
-    - 退出: lifespan finally 里 registry.dispose_all()(Windows 文件锁,
-      不关则 tmp 目录删不掉);随后置空 app.state.registry(关残留窗——
-      不让已 dispose 的测试 Registry 被后续无夹具的 lifespan 误用)并
-      清理 tmp 目录(删不掉告警留痕,不静默吞)。
+
+def build_app(seed_dir, data_dir):
+    """app 工厂(Task 7): 建 Registry → 预置 app.state.registry → 返回 app。
+
+    Host 中间件已在 app.main 模块级挂好(add_middleware,经 scope["app"]
+    延迟取 registry),工厂只负责把 Registry 放进 app.state——lifespan 见
+    已挂则沿用(测试注入缝),不再加载仓库 seed。调用方用完必须复位
+    app.state.registry = None(见 app_test_client 的 finally)。
    """
-    from fastapi.testclient import TestClient
-
-    registry = TenantRegistry.load_from_dir(SEED_DIR, data_dir=tmp_path)
+    registry = TenantRegistry.load_from_dir(seed_dir, data_dir=data_dir)
    app.state.registry = registry
+    return app
+
+
+def _rmtree_warn(path):
    try:
-        with TestClient(app) as c:
+        shutil.rmtree(path)
+    except OSError as exc:
+        warnings.warn(f"测试夹具清理 tmp 目录失败(文件锁?): {exc}")
+        shutil.rmtree(path, ignore_errors=True)
+
+
+class TenantTestClient(TestClient):
+    """WS 连接默认 Host 取 base_url 域名。
+
+    starlette 的 websocket_connect 硬编码 ``ws://testserver`` 拼绝对 URL,
+    base_url 对 WS 不生效——不补这层,既有 WS 测试的 Host 全落 testserver
+    (未知域名,被中间件 1008 拒)。显式传 headers={"host": ...} 仍可覆盖
+    (切租户用,见 test_host_routing)。"""
+
+    def __init__(self, *args, ws_default_host: str = "", **kwargs):
+        super().__init__(*args, **kwargs)
+        self._ws_default_host = ws_default_host
+
+    def websocket_connect(self, url, subprotocols=None, **kwargs):
+        headers = dict(kwargs.pop("headers", None) or {})
+        if self._ws_default_host:
+            headers.setdefault("host", self._ws_default_host)
+        kwargs["headers"] = headers
+        return super().websocket_connect(url, subprotocols=subprotocols, **kwargs)
+
+
+@contextmanager
+def app_test_client(app_instance, base_url, cleanup_root):
+    """TestClient 包装:统一退出纪律——lifespan finally 已 dispose_all
+    (Windows 文件锁),再置空 app.state.registry(关残留窗:不让已 dispose
+    的测试 Registry 被后续无夹具的 lifespan 误用),最后清 tmp 目录
+    (删不掉告警留痕,不静默吞)。"""
+    try:
+        with TenantTestClient(
+                app_instance, base_url=base_url,
+                ws_default_host=base_url.removeprefix("http://")) as c:
            yield c
    finally:
-        app.state.registry = None
-        try:
-            shutil.rmtree(tmp_path)
-        except OSError as exc:
-            warnings.warn(f"client 夹具清理 tmp 目录失败(文件锁?): {exc}")
-            shutil.rmtree(tmp_path, ignore_errors=True)
+        app_instance.state.registry = None
+        _rmtree_warn(cleanup_root)
+
+
+@pytest.fixture
+def app_client(tmp_path):
+    """TestClient(base_url=jiancai 域名) + 每测试独立 Registry(tmp data_dir)。
+
+    - Host 中间件按请求 Host 逐请求 resolve_site:engine 懒建在 TestClient
+      的 portal 事件循环里(连接/循环绑定安全)。测试线程需要直接造数时
+      不得用该 engine——asyncio.run 另起循环,跨循环复用常驻连接会死锁;
+      直连断言走 async_db,或经真实路由(WS/REST)。
+    - 切换租户: 请求级 headers={"host": ...} 覆盖(WS 同理),不用另开
+      TestClient——同一 app 上两个 client 各跑一次 lifespan 会共享
+      registry 而 engine 各绑一个循环,跨循环死锁。
+    """
+    build_app(SEED_DIR, tmp_path)
+    with app_test_client(app, f"http://{JIANCAI_HOST}", tmp_path) as c:
+        yield c
+
+
+@pytest.fixture
+def client(app_client):
+    """既有测试的 client 夹具(Task 7 起基于 app_client,仅换名)。"""
+    yield app_client
@@ -0,0 +1,141 @@
+# Task 7: Host 中间件租户路由——纯 ASGI(http+websocket 双覆盖)。
+#
+# 本文件的存在意义(前序审查实证过的坑): 若中间件走了 BaseHTTPMiddleware
+# (@app.middleware("http")),WS 请求会整个绕过租户解析落 get_runtime 兜底,
+# auto 域名的 WS 流量静默写进 jiancai 库(串库不崩溃,极难发现)。
+# test_ws_routes_to_correct_tenant_db 用双库断言钉死这条红线。
+
+import json
+import shutil
+
+import pytest
+from starlette.websockets import WebSocketDisconnect
+
+from tests.conftest import (
+    AUTO_HOST,
+    JIANCAI_HOST,
+    SEED_DIR,
+    app_test_client,
+    build_app,
+)
+
+NOPE_HOST = "nope.example.com"
+
+
+# ---- 平台级豁免 -------------------------------------------------------------
+
+
+def test_known_host_serves_health(app_client):
+    # /health 在任何 Host 下可达(部署健康检查 curl 127.0.0.1:8018,
+    # Host 是 IP:端口,对 registry 必然未知)
+    assert app_client.get("/health").status_code == 200
+    r = app_client.get("/health", headers={"host": NOPE_HOST})
+    assert r.status_code == 200
+
+
+def test_health_exemption_is_exact_match(app_client):
+    # 精确 == "/health",/healthxxx 不豁免:未知 Host 由中间件(而非路由)
+    # 拒绝——text/plain 的 PlainTextResponse 区别于路由 404 的 JSON body
+    r = app_client.get("/healthxxx", headers={"host": NOPE_HOST})
+    assert r.status_code == 404
+    assert r.headers["content-type"].startswith("text/plain")
+
+
+# ---- 未知 Host --------------------------------------------------------------
+
+
+def test_unknown_host_gets_404(app_client):
+    r = app_client.get("/api/chat/session", headers={"host": NOPE_HOST})
+    assert r.status_code == 404
+
+
+def test_unknown_host_ws_rejected_with_1008(app_client):
+    # 探针实证: 服务器未 accept 直接 close 时,TestClient 在 __enter__ 即抛
+    # WebSocketDisconnect,code/reason 保留——这就是断言方式
+    with pytest.raises(WebSocketDisconnect) as ei:
+        with app_client.websocket_connect("/ws/chat", headers={"host": NOPE_HOST}):
+            pass  # pragma: no cover — 未 accept 即 close,进不来
+    assert ei.value.code == 1008
+
+
+# ---- 停用租户 ---------------------------------------------------------------
+
+
+@pytest.fixture
+def disabled_client(tmp_path):
+    """jiancai 停用版 app_client: 拷 seed 到 tmp、改 status 再加载。
+
+    加载前改盘(而非 registry_reload): 无需跨事件循环碰 registry 锁,
+    最简;reload 语义本身已被 Task 5a 的测试钉过,此处不重复走。"""
+    seed = tmp_path / "seed"
+    shutil.copytree(SEED_DIR, seed)
+    tenant_json = seed / "jiancai" / "tenant.json"
+    data = json.loads(tenant_json.read_text(encoding="utf-8"))
+    data["status"] = "disabled"
+    tenant_json.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")
+    a = build_app(seed, tmp_path / "data")
+    with app_test_client(a, f"http://{JIANCAI_HOST}", tmp_path) as c:
+        yield c
+
+
+def test_disabled_tenant_gets_pause_page(disabled_client):
+    r = disabled_client.get("/api/chat/session")
+    assert r.status_code == 503
+    assert "暂停" in r.text
+    # 中性页: 不暴露租户 slug/域名
+    assert "jiancai" not in r.text.lower()
+    assert JIANCAI_HOST not in r.text
+
+
+def test_disabled_tenant_ws_rejected_with_1013(disabled_client):
+    with pytest.raises(WebSocketDisconnect) as ei:
+        with disabled_client.websocket_connect("/ws/chat"):
+            pass  # pragma: no cover
+    assert ei.value.code == 1013
+
+
+# ---- 租户正确性(防串库红线) --------------------------------------------------
+
+
+def test_default_host_routes_to_jiancai(app_client):
+    # app_client 不带 host 头 → base_url 的 jiancai 域(既有全量测试的语义)
+    r = app_client.post("/api/chat/session", json={})
+    assert r.status_code == 200
+
+
+def test_ws_routes_to_correct_tenant_db(app_client, tmp_path):
+    """WS 经 auto 域名连接 → 会话必须落在 auto 租户库。
+
+    三层断言:
+    1. 文件级——此刻只碰过 auto(WS 用显式 host;lifespan 清扫因无 db
+       文件而跳过),auto.db 已懒建而 jiancai.db 从未创建;
+    2. REST resume 经 auto 域名命中同一 sid(is_new=False 且回原 sid);
+    3. 同一 sid 经 jiancai 域名(base_url 默认)不可见——resume miss 落
+       is_new=True(新建了另一个会话),证明 jiancai 库里没有它。
+    (经 /api/chat/session 的 resume 探测而非 /messages: 后者只放行
+    app-source 会话,WS 造的 standalone 会话两库都会 404,区分不了;
+    resume 探测不触发 AI 调用。)"""
+    with app_client.websocket_connect(
+            "/ws/chat", headers={"host": AUTO_HOST}) as ws:
+        ws.send_text(json.dumps({"source": "standalone"}))
+        init = ws.receive_json()
+        assert init["type"] == "init"
+        sid = init["session_id"]
+
+    # 1. 文件级(jiancai 还没被任何请求 resolve 过)
+    assert (tmp_path / "auto.db").exists()
+    assert not (tmp_path / "jiancai.db").exists()
+
+    # 2. auto 库命中
+    r_auto = app_client.post(
+        "/api/chat/session", json={"session_id": sid},
+        headers={"host": AUTO_HOST})
+    assert r_auto.status_code == 200
+    assert r_auto.json()["is_new"] is False
+    assert r_auto.json()["session_id"] == sid
+
+    # 3. jiancai 库没有(miss → 新建,is_new=True、sid 不同)
+    r_jc = app_client.post("/api/chat/session", json={"session_id": sid})
+    assert r_jc.status_code == 200
+    assert r_jc.json()["is_new"] is True
+    assert r_jc.json()["session_id"] != sid