ai-kefu-admin/tests/conftest.py

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.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
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


# 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)。
    """
    registry = TenantRegistry.load_from_dir(seed_dir, data_dir=data_dir)
    app.state.registry = registry
    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
    (Windows 文件锁),再置空 app.state.registry(关残留窗:不让已 dispose
    的测试 Registry 被后续无夹具的 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
        _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