""" Shared pytest fixtures for integration tests that exercise the live API service. """ from __future__ import annotations import asyncio import json import os import subprocess import sys import uuid from collections.abc import AsyncGenerator from pathlib import Path import anyio import httpx import pytest import pytest_asyncio from dotenv import load_dotenv PROJECT_ROOT = Path(__file__).resolve().parents[2] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) load_dotenv(PROJECT_ROOT / ".env", override=False) load_dotenv(PROJECT_ROOT / "test/.env.test", override=False) API_BASE_URL = os.getenv("TEST_BASE_URL", "http://localhost:5050").rstrip("/") ADMIN_LOGIN = os.getenv("TEST_USERNAME") ADMIN_PASSWORD = os.getenv("TEST_PASSWORD") _ADMIN_TOKEN_CACHE: str | None = None # 防止并发首次登录造成雪崩（触发 429 限流 / 423 账户锁定） _login_lock = asyncio.Lock() HTTP_TIMEOUT = httpx.Timeout(60.0, connect=5.0) SANDBOX_CONTAINER_PREFIX = os.getenv("YUXI_SANDBOX_CONTAINER_PREFIX", "yuxi-sandbox") @pytest.fixture(scope="session", autouse=True) def ensure_live_api_schema(): if not ADMIN_LOGIN or not ADMIN_PASSWORD: return async def run_schema_setup() -> None: from yuxi.storage.postgres.manager import pg_manager pg_manager.initialize() await pg_manager.create_tables() await pg_manager.ensure_business_schema() await pg_manager.ensure_knowledge_schema() await pg_manager.ensure_external_schema() await pg_manager.ensure_scheduler_schema() await pg_manager.ensure_channel_schema() anyio.run(run_schema_setup) def _require_admin_credentials() -> tuple[str, str]: if not ADMIN_LOGIN or not ADMIN_PASSWORD: pytest.skip("Integration credentials are not configured via TEST_USERNAME / TEST_PASSWORD.") return ADMIN_LOGIN, ADMIN_PASSWORD @pytest_asyncio.fixture(scope="function") async def test_client() -> AsyncGenerator[httpx.AsyncClient, None]: async with httpx.AsyncClient(base_url=API_BASE_URL, timeout=HTTP_TIMEOUT, follow_redirects=True) as client: yield client @pytest_asyncio.fixture(scope="function") async def admin_token() -> str: """获取管理员 token，带缓存预检与限流退避重试。 设计要点： - 缓存命中时先用 GET /api/auth/me 校验 token 有效性（GET 不触发登录限流）， 避免每个用例都重新登录造成雪崩。 - token 失效才重新登录；登录遇到 429 按 Retry-After 退避重试，最多 2 次。 - 模块级 _login_lock 防止多用例并发首次登录。 """ global _ADMIN_TOKEN_CACHE async with httpx.AsyncClient( base_url=API_BASE_URL, timeout=HTTP_TIMEOUT, follow_redirects=True, ) as client: # [1] 预检：缓存命中时用 GET /api/auth/me 校验 token 是否仍然有效 if _ADMIN_TOKEN_CACHE: try: me_resp = await client.get( "/api/auth/me", headers={"Authorization": f"Bearer {_ADMIN_TOKEN_CACHE}"}, ) except httpx.HTTPError as exc: pytest.fail(f"Failed to verify admin token via /api/auth/me: {exc}") if me_resp.status_code == 200: return _ADMIN_TOKEN_CACHE # type: ignore[return-value] if me_resp.status_code == 423: pytest.fail( "Admin account is login-locked. Unlock it in DB (reset login_failed_count / " "login_locked_until) and restart api-dev before re-running tests." ) # 401 / 其他状态：token 已失效，清缓存走重新登录流程 _ADMIN_TOKEN_CACHE = None # [2] 登录：缓存失效或首次。加锁防止并发首次登录雪崩 username, password = _require_admin_credentials() async with _login_lock: # 双检：持锁后再次确认缓存（其他用例可能已登录成功） if _ADMIN_TOKEN_CACHE: return _ADMIN_TOKEN_CACHE # type: ignore[return-value] response = await _do_login_with_backoff(client, username, password) if response.status_code == 401: first_run_response = await client.get("/api/auth/check-first-run") if first_run_response.status_code == 200 and first_run_response.json().get("first_run", False): pytest.fail( "Super admin account has not been initialized. Complete `/api/auth/initialize` before " "running integration tests." ) if response.status_code != 200: pytest.fail(f"Failed to authenticate as admin (status={response.status_code}): {response.text}") token = response.json().get("access_token") if not token: pytest.fail("Admin authentication did not return an access token.") _ADMIN_TOKEN_CACHE = token return token async def _do_login_with_backoff( client: httpx.AsyncClient, username: str, password: str, ) -> httpx.Response: """执行登录，遇到 429 按 Retry-After 退避重试，最多 2 次。""" max_retries = 2 response = await client.post( "/api/auth/token", data={"username": username, "password": password}, ) for attempt in range(max_retries): if response.status_code != 429: break # 解析 Retry-After（服务端返回的 details.retry_after 优先，header 次之） retry_after: float = 5.0 try: body = response.json() retry_after = float(body.get("error", {}).get("details", {}).get("retry_after", retry_after)) except (ValueError, KeyError): header_val = response.headers.get("Retry-After") if header_val: try: retry_after = float(header_val) except ValueError: pass await anyio.sleep(retry_after + 1.0) response = await client.post( "/api/auth/token", data={"username": username, "password": password}, ) return response @pytest.fixture(scope="function") def admin_headers(admin_token: str) -> dict[str, str]: return {"Authorization": f"Bearer {admin_token}"} @pytest.fixture(scope="session", autouse=True) def cleanup_test_knowledge_databases(): async def run_cleanup() -> None: global _ADMIN_TOKEN_CACHE if not ADMIN_LOGIN or not ADMIN_PASSWORD: return if not _ADMIN_TOKEN_CACHE: async with httpx.AsyncClient( base_url=API_BASE_URL, timeout=HTTP_TIMEOUT, follow_redirects=True, ) as bootstrap_client: response = await bootstrap_client.post( "/api/auth/token", data={"username": ADMIN_LOGIN, "password": ADMIN_PASSWORD}, ) if response.status_code != 200: return token = response.json().get("access_token") if not token: return _ADMIN_TOKEN_CACHE = token headers = {"Authorization": f"Bearer {_ADMIN_TOKEN_CACHE}"} async with httpx.AsyncClient(base_url=API_BASE_URL, timeout=HTTP_TIMEOUT, follow_redirects=True) as client: try: list_response = await client.get("/api/knowledge/databases", headers=headers) except Exception as exc: print(f"Warning: Failed to list knowledge databases for cleanup: {exc}") return if list_response.status_code != 200: return databases = list_response.json().get("databases", []) prefixes = ("pytest_", "py_test") for entry in databases: name = entry.get("name") or "" slug = entry.get("slug") if not slug or not isinstance(name, str) or not name.startswith(prefixes): continue try: delete_response = await client.delete(f"/api/knowledge/databases/{slug}", headers=headers) if delete_response.status_code not in (200, 404): print(f"Warning: Failed to cleanup knowledge database {slug}: {delete_response.text}") except Exception as exc: print(f"Warning: Exception during cleanup of {slug}: {exc}") try: anyio.run(run_cleanup) except Exception as exc: print(f"Warning: Exception during session cleanup startup: {exc}") yield try: anyio.run(run_cleanup) except Exception as exc: print(f"Warning: Exception during session cleanup teardown: {exc}") def _docker_api_request(method: str, path: str) -> list[dict] | dict: cmd = [ "curl", "-sS", "--unix-socket", os.getenv("YUXI_DOCKER_API_SOCKET", "/var/run/docker.sock"), "-X", method, f"{os.getenv('YUXI_DOCKER_API_BASE', 'http://localhost').rstrip('/')}{path}", ] result = subprocess.run(cmd, capture_output=True, text=True, timeout=15, check=False) if result.returncode != 0: raise RuntimeError(result.stderr.strip() or "Docker API request failed") text = (result.stdout or "").strip() if not text: return {} return json.loads(text) def _cleanup_sandbox_containers() -> None: try: containers = _docker_api_request("GET", "/containers/json?all=true") except Exception as exc: print(f"Warning: Failed to list sandbox containers for cleanup: {exc}") return for container in containers if isinstance(containers, list) else []: names = container.get("Names") or [] if not any(name.lstrip("/").startswith(f"{SANDBOX_CONTAINER_PREFIX}-") for name in names): continue container_id = container.get("Id") if not container_id: continue try: _docker_api_request("POST", f"/containers/{container_id}/stop?t=2") except Exception: pass try: _docker_api_request("DELETE", f"/containers/{container_id}?force=true") except Exception as exc: print(f"Warning: Failed to cleanup sandbox container {container_id[:12]}: {exc}") @pytest.fixture(scope="session", autouse=True) def cleanup_test_sandboxes(): _cleanup_sandbox_containers() yield _cleanup_sandbox_containers() @pytest_asyncio.fixture(scope="function") async def standard_user(test_client: httpx.AsyncClient, admin_headers: dict[str, str]) -> AsyncGenerator[dict, None]: username = f"pytest_user_{uuid.uuid4().hex[:8]}" password = f"Pw!{uuid.uuid4().hex[:8]}" # 用户隔离重构后所有登录用户必须绑定部门，创建时显式指定一个已存在部门 dept_response = await test_client.get("/api/departments", headers=admin_headers) if dept_response.status_code != 200 or not dept_response.json(): pytest.fail(f"No department available to bind standard user: {dept_response.text}") department_id = dept_response.json()[0]["id"] response = await test_client.post( "/api/auth/users", json={"username": username, "password": password, "role": "user", "department_id": department_id}, headers=admin_headers, ) if response.status_code != 200: pytest.fail(f"Failed to create standard user (status={response.status_code}): {response.text}") user_payload = response.json() login_response = await test_client.post( "/api/auth/token", data={"username": user_payload["uid"], "password": password}, ) if login_response.status_code != 200: pytest.fail( f"Failed to authenticate as standard user (status={login_response.status_code}): {login_response.text}" ) access_token = login_response.json().get("access_token") if not access_token: pytest.fail("Standard user login succeeded but no access token was returned.") try: yield { "user": user_payload, "password": password, "headers": {"Authorization": f"Bearer {access_token}"}, } finally: cleanup_error = None for _ in range(3): response = await test_client.delete(f"/api/auth/users/{user_payload['id']}", headers=admin_headers) if response.status_code in (200, 404): cleanup_error = None break cleanup_error = response await anyio.sleep(0.3) if cleanup_error is not None: assert cleanup_error.status_code == 200, ( f"Failed to cleanup test user {user_payload['uid']}: {cleanup_error.text}" ) @pytest_asyncio.fixture(scope="function") async def knowledge_database( test_client: httpx.AsyncClient, admin_headers: dict[str, str], ) -> AsyncGenerator[dict, None]: import time unique_id = uuid.uuid4().hex timestamp = int(time.time() * 1000000) db_name = f"pytest_kb_{timestamp}_{unique_id}" kb_id = None try: create_response = await test_client.post( "/api/knowledge/databases", json={ "database_name": db_name, "description": "Pytest managed knowledge base", "embedding_model_spec": "siliconflow-cn:Pro/BAAI/bge-m3", "kb_type": "milvus", "additional_params": {}, }, headers=admin_headers, ) if create_response.status_code == 200: db_payload = create_response.json() kb_id = db_payload["kb_id"] elif create_response.status_code == 409: error_detail = create_response.json().get("detail", "") pytest.fail(f"Knowledge database name conflict: {error_detail}. Please clean up old test databases first.") else: pytest.fail( f"Failed to create knowledge database (status={create_response.status_code}): {create_response.text}" ) yield db_payload finally: if kb_id: try: delete_response = await test_client.delete(f"/api/knowledge/databases/{kb_id}", headers=admin_headers) if delete_response.status_code != 200: print(f"Warning: Failed to cleanup knowledge database {kb_id}: {delete_response.text}") except Exception as exc: print(f"Warning: Exception during cleanup of {kb_id}: {exc}")