vent_agent_2.0/tests/conftest.py

"""pytest 全局装置。

**必须在导入任何项目模块之前设置环境变量**：``config`` 在导入时就读
``.env`` 与环境变量并定稿，``db.db_conn`` 在导入时就建好连接适配器。
所以下面的 ``os.environ`` 赋值必须在 ``import config`` 之前发生——
这也是把它们放在文件最顶部、并设 ``DEEP_AGENT_SKIP_DOTENV`` 的原因：
不跳过 .env 的话，开发者本机的 .env 会覆盖这里的测试配置，
同一套测试在不同机器上跑出不同结果。
"""

from __future__ import annotations

import os
import shutil
import tempfile
import uuid
from pathlib import Path

# ---------- 测试环境（必须在项目导入前设置）----------
# 用 setdefault 而非直接赋值：外部已给出的 DB_DIALECT / DB_* 优先生效，
# 这样同一套测试可以直接跑在 MySQL 上验证方言无关性——
#   DB_DIALECT=mysql DB_HOST=127.0.0.1 DB_PORT=3399 DB_USER=root DB_NAME=t python -m pytest
_TMP_ROOT = Path(tempfile.mkdtemp(prefix="deep_agent_test_"))
os.environ["DEEP_AGENT_ENV_PRIORITY"] = "none"   # 完全忽略本机 .env
os.environ.setdefault("DB_DIALECT", "sqlite")
os.environ.setdefault("DB_SQLITE_PATH", str(_TMP_ROOT / "test.sqlite"))
os.environ.setdefault("CHECKPOINTER", "memory")   # 测试不需要落盘图状态
os.environ["CTX_COMPACT_ENABLED"] = "0"           # 不引入摘要中间件，测试更纯粹
os.environ["CONTEXT_DEMO_ENABLED"] = "0"          # 示例扩展由专门用例启用
os.environ["OPENAI_API_KEY"] = "sk-test-key"
os.environ["OPENAI_BASE_URL"] = "https://example.invalid"
os.environ["TFAGENTS_DAILY_TOKEN_LIMIT"] = "0"
os.environ["BILLING_ALLOW_UNPRICED"] = "0"
os.environ["LOG_LEVEL"] = "WARNING"

import pytest  # noqa: E402

import config  # noqa: E402
from db import db_conn  # noqa: E402


@pytest.fixture(scope="session", autouse=True)
def _schema() -> None:
    """整个测试会话共用一份建好的表结构。"""
    db_conn.init_schema()


@pytest.fixture(autouse=True)
def _clean_tables():
    """每个用例前清空业务表，避免用例间互相污染。

    删表数据而不是重建 schema：重建的代价在 sqlite 上也不小，而这里
    只需要一个干净的状态。
    """
    for table in (
        "messages",
        "execution_runs",
        "request_keys",
        "billing_pending",
        "sessions",
        "session_context_usage",
        "context_archives",
        "billing_usage",
        "billing_pricing",
        "daily_quota",
        "user_preferences",
        "schedule_runs",
        "schedules",
    ):
        db_conn.execute(f"DELETE FROM {table}", ())
    # 工具注册表是模块级单例，清空后重新装载，避免用例顺序影响工具清单
    from tools import registry

    registry.reset_for_tests()
    registry.load_builtin_tools()
    yield


@pytest.fixture(scope="session")
def dialect() -> str:
    """当前测试跑的方言。用于需要区分方言的断言。"""
    return config.DB_DIALECT


@pytest.fixture
def pricing_ready():
    """确保单价表里有当前模型的条目，供计费相关用例使用。"""
    from billing import pricing

    pricing.seed_default_pricing()
    return pricing


@pytest.fixture
def session_id() -> str:
    """建一个测试会话，返回 session_id。"""
    from db import chat_store

    return chat_store.create_session("tester", "full", "测试会话")["id"]


@pytest.fixture
def unique_user() -> str:
    """每次调用返回一个新用户名，避免额度/流水跨用例串味。"""
    return f"u_{uuid.uuid4().hex[:8]}"


def pytest_sessionfinish(session, exitstatus) -> None:
    """收尾清理临时目录。"""
    shutil.rmtree(_TMP_ROOT, ignore_errors=True)