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

Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
浅念
2026-10-12 20:30:30 +08:00
parent 1824bb1590
commit eaa5bf7e51
5 changed files with 338 additions and 50 deletions
+80 -24
View File
@@ -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