Files
ai-kefu-admin/tests/conftest.py
T

234 lines
8.8 KiB
Python

import asyncio
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
from urllib.parse import urlsplit
from app.engine.models import Base
# Task 5b 装配切换: app 不再有全局 engine/async_session——旧 conftest 里
# "先 import app.main 再 import chat_rest"的循环导入次序坑随全局单例一并
# 消失(routes 不再反向 import app.main),直接导入即可。
from app.main import app
from app.platform.db import make_platform_engine
from app.platform.registry import TenantRegistry
REPO_ROOT = Path(__file__).resolve().parents[1]
SEED_DIR = REPO_ROOT / "seed_profiles"
async def _create_all(engine):
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
@pytest.fixture
def async_db():
"""Fresh temp-file SQLite sessionmaker for direct async-DB tests
(e.g. testing ConversationManager without HTTP). Isolated per test.
刻意不走 Registry 引擎: 本夹具被各测试反复 asyncio.run(每次全新事件
循环),NullPool 连接即借即还才能跨循环安全;Registry 的 pool=2 常驻
连接与创建它的循环绑定,只能活在单一事件循环里(client 夹具的
TestClient 循环,见下)。
"""
fd, path = tempfile.mkstemp(suffix=".db")
os.close(fd)
engine = create_async_engine(
f"sqlite+aiosqlite:///{path}", poolclass=NullPool
)
asyncio.run(_create_all(engine))
sm = async_sessionmaker(engine, expire_on_commit=False)
yield sm
asyncio.run(engine.dispose())
try:
os.remove(path)
except PermissionError:
pass
@pytest.fixture(autouse=True)
def _reset_rate_limit():
"""Reset the in-memory rate limiter before and after every test so that
module-global bucket state can never leak across tests."""
from app.engine.services.rate_limit import _reset_for_test
_reset_for_test()
yield
_reset_for_test()
@pytest.fixture(autouse=True)
def _reset_conn_limiter():
"""Task 8: 同上——WS 连接限流器的模块级单例计数前后复位,跨测试不泄漏
(TestClient 所有连接同落 ip="testclient",不复位会互相顶名额)。"""
from app.engine.services.conn_limiter import _reset_for_test
_reset_for_test()
yield
_reset_for_test()
@pytest.fixture(autouse=True)
def _reset_platform_auth():
"""Phase 2 Task 2: 平台认证模块级单例(内存 token 表/登录限流失败
计数)前后复位,跨测试不泄漏——token/限流窗不隔离会串测试(前一脚的
锁定状态顶掉后一脚的登录);注入的 TTL/窗长一并还原默认。"""
from app.platform.auth import _reset_for_test
_reset_for_test()
yield
_reset_for_test()
@pytest.fixture
def jiancai_runtime():
"""直连测试用的 jiancai TenantRuntime(Task 6: manager 的 runtime 必填)。
委托 build_jiancai_runtime()(见其注释: 只走加载期,无 engine 无句柄,
跨测试共享单例安全)。
"""
return build_jiancai_runtime()
_JIANCAI_RUNTIME = None
def build_jiancai_runtime():
"""(非夹具)jiancai TenantRuntime——给批量集成测试的 _setup 辅助线用,
免去百处测试签名逐一加夹具。
只走加载期(config/prompts/compiled_wall 就位),从不 ensure_ready——
engine/session_factory 保持 None: 直连测试自管 async_db(NullPool 跨
循环纪律;Task 9 起 lead_push_service 无模块级 session 兜底,session_factory
由调用方/测试显式传)。无 engine 即无句柄,无需清理;
data_dir 仅占位(只拼 db_path 字符串,无任何 IO——从不 ensure_ready,
不会创建 tenants/ 目录或库文件),模块级缓存单例跨测试复用。
"""
global _JIANCAI_RUNTIME
if _JIANCAI_RUNTIME is None:
reg = TenantRegistry.load_from_dir(
SEED_DIR, data_dir=REPO_ROOT / "tenants")
_JIANCAI_RUNTIME = next(
rt for rt in reg.runtimes() if rt.config.slug == "jiancai")
return _JIANCAI_RUNTIME
_AUTO_RUNTIME = None
def build_auto_runtime():
"""(非夹具)auto TenantRuntime——C-1 终审回归锁用(两租户墙差异断言),
纪律同 build_jiancai_runtime: 只走加载期,无 engine 无句柄,
模块级缓存单例跨测试复用。"""
global _AUTO_RUNTIME
if _AUTO_RUNTIME is None:
reg = TenantRegistry.load_from_dir(
SEED_DIR, data_dir=REPO_ROOT / "tenants")
_AUTO_RUNTIME = next(
rt for rt in reg.runtimes() if rt.config.slug == "auto")
return _AUTO_RUNTIME
# Task 7: app_client 默认域名——既有测试不带 host 头,默认落 jiancai
# (与 5b 之前 default_runtime 兜底的语义一致,既有断言零改动)
JIANCAI_HOST = "chat2.agent.hnxinyouhe.com"
AUTO_HOST = "chat.agent.hnxinyouhe.com"
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)。
Phase 2 Task 1 起同时预置 platform_engine(data_dir 下的临时平台库,
同一注入缝)——不预置则 lifespan 会在仓库根建出真实 platform.db;
schema+bootstrap 由 lifespan 统一跑(幂等),dispose 也在 lifespan
finally(与 registry.dispose_all 同处)。
"""
registry = TenantRegistry.load_from_dir(seed_dir, data_dir=data_dir)
app.state.registry = registry
# data_dir 可能尚不存在(租户库由 _ensure_ready 按需 mkdir,平台库由
# lifespan 立即连接——这里先建目录,否则 sqlite "unable to open database
# file")
Path(data_dir).mkdir(parents=True, exist_ok=True)
app.state.platform_engine = make_platform_engine(
Path(data_dir) / "platform.db")
return app
def _rmtree_warn(path):
try:
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 +
platform_engine.dispose(Windows 文件锁),再置空 app.state.registry /
platform_engine(关残留窗:不让已 dispose 的测试对象被后续无夹具的
lifespan 误用),最后清 tmp 目录(删不掉告警留痕,不静默吞)。"""
try:
with TenantTestClient(
app_instance, base_url=base_url,
ws_default_host=urlsplit(base_url).netloc) as c:
yield c
finally:
app_instance.state.registry = None
app_instance.state.platform_engine = 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 各绑一个循环,跨循环死锁。
"""
a = build_app(SEED_DIR, tmp_path)
with app_test_client(a, f"http://{JIANCAI_HOST}", tmp_path) as c:
yield c
@pytest.fixture
def client(app_client):
"""既有测试的 client 夹具(Task 7 起基于 app_client,仅换名)。"""
yield app_client