bid-doctor/server/tests/conftest.py

"""接口集成测试：跑在真实的 PostgreSQL 上（行级安全策略生效），Valkey 用第 1 号库。

需要先起依赖：docker compose -f deploy/compose/compose.yaml up -d postgres valkey gotenberg
每次测试会话重建 bid_doctor_test 库并跑全部迁移。
"""

from __future__ import annotations

import itertools
import os
import tempfile
from collections.abc import AsyncIterator
from pathlib import Path
from typing import Any

PG = os.environ.get("BD_TEST_PG", "127.0.0.1:54329")
os.environ.update(
    {
        "BD_ENV": "test",
        "BD_DATABASE_URL": f"postgresql+psycopg://bd_app:dev-app@{PG}/bid_doctor_test",
        "BD_MIGRATE_DATABASE_URL": f"postgresql+psycopg://bd_owner:dev-owner@{PG}/bid_doctor_test",
        "BD_VALKEY_URL": os.environ.get("BD_TEST_VALKEY", "redis://127.0.0.1:63799/1"),
        "BD_STORAGE_DIR": tempfile.mkdtemp(prefix="bd-test-objects-"),
        "BD_AI_PROVIDER": "fake",
        "BD_SMS_PROVIDER": "dev",
    }
)

import httpx  # noqa: E402
import psycopg  # noqa: E402
import pytest  # noqa: E402

FIXTURES = Path(__file__).parent / "fixtures"
_phones = itertools.count(13900000001)


def _recreate_database() -> None:
    with psycopg.connect(f"postgresql://bd_owner:dev-owner@{PG}/postgres", autocommit=True) as conn:
        conn.execute("DROP DATABASE IF EXISTS bid_doctor_test WITH (FORCE)")
        conn.execute("CREATE DATABASE bid_doctor_test")
        conn.execute("GRANT CONNECT ON DATABASE bid_doctor_test TO bd_app")


@pytest.fixture(scope="session", autouse=True)
def database() -> None:
    _recreate_database()
    from alembic import command
    from alembic.config import Config

    command.upgrade(Config(str(Path(__file__).resolve().parents[1] / "alembic.ini")), "head")


@pytest.fixture(scope="session", autouse=True)
async def _valkey() -> AsyncIterator[None]:
    from bid_doctor.core.valkey import valkey

    await valkey().flushdb()
    yield
    await valkey().flushdb()


@pytest.fixture(scope="session")
async def app_client() -> AsyncIterator[httpx.AsyncClient]:
    from bid_doctor.app import create_app

    transport = httpx.ASGITransport(app=create_app())
    async with httpx.AsyncClient(transport=transport, base_url="http://test/api/v1") as c:
        yield c


class User:
    """一个注册好的企业用户（桌面版 Bearer 令牌）。"""

    def __init__(self, client: httpx.AsyncClient, token: str, me: dict[str, Any], phone: str) -> None:
        self.client, self.token, self.me, self.phone = client, token, me, phone
        self.headers = {"Authorization": f"Bearer {token}"}

    async def get(self, path: str, **kw: Any) -> httpx.Response:
        return await self.client.get(path, headers=self.headers, **kw)

    async def post(self, path: str, **kw: Any) -> httpx.Response:
        return await self.client.post(path, headers=self.headers, **kw)

    async def put(self, path: str, **kw: Any) -> httpx.Response:
        return await self.client.put(path, headers=self.headers, **kw)

    @property
    def tenant_id(self) -> str:
        return str(self.me["tenant"]["id"])


async def register(client: httpx.AsyncClient, name: str = "测试集成公司") -> User:
    phone = str(next(_phones))
    assert (await client.post("/auth/sms-codes", json={"phone": phone, "purpose": "register"})).status_code == 202
    r = await client.post(
        "/auth/register",
        json={
            "phone": phone,
            "sms_code": "000000",
            "tenant_name": name,
            "display_name": "张工",
            "terms_version": "v1",
            "client_kind": "desktop",
        },
    )
    assert r.status_code == 201, r.text
    body = r.json()
    return User(client, body["token"], body["me"], phone)


@pytest.fixture
async def user(app_client: httpx.AsyncClient) -> User:
    return await register(app_client)


async def upload(u: User, filename: str, data: bytes) -> httpx.Response:
    r = await u.post("/uploads", json={"filename": filename, "size_bytes": len(data), "purpose": "project_document"})
    assert r.status_code == 201, r.text
    up = r.json()
    size = up["part_size"]
    for n in range(1, up["part_count"] + 1):
        assert (
            await u.put(f"/uploads/{up['id']}/parts/{n}", content=data[(n - 1) * size : n * size])
        ).status_code == 204
    return await u.post(f"/uploads/{up['id']}/complete")


async def project_with_tender(u: User, fixture: str = "tender-mini.pdf") -> tuple[str, str]:
    """建项目并挂上一份招标文件，返回 (项目编号, 文档编号)。"""
    project = (await u.post("/projects", json={})).json()
    done = await upload(u, fixture, (FIXTURES / fixture).read_bytes())
    assert done.status_code == 200, done.text
    r = await u.post(f"/projects/{project['id']}/documents", json={"upload_id": done.json()["id"], "role": "tender"})
    assert r.status_code == 201, r.text
    return project["id"], r.json()["id"]


async def run_parse(u: User, document_id: str) -> dict[str, Any]:
    """发起解析，并在测试进程里直接跑流水线（不经过 worker）。"""
    from uuid import UUID

    from bid_doctor.document.pipeline import runner

    r = await u.post(f"/documents/{document_id}/parse-jobs", json={"expected_charge": 1})
    assert r.status_code == 201, r.text
    await runner.run(UUID(u.tenant_id), UUID(r.json()["id"]))
    return (await u.get(f"/parse-jobs/{r.json()['id']}")).json()


async def balance(u: User) -> int:
    accounts = (await u.get("/quota")).json()["accounts"]
    return next(a["balance"] for a in accounts if a["unit"] == "parse")