""" Shared pytest fixtures for integration tests that exercise the live API service. """ from __future__ import annotations 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 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() 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: global _ADMIN_TOKEN_CACHE if _ADMIN_TOKEN_CACHE: return _ADMIN_TOKEN_CACHE username, password = _require_admin_credentials() 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": username, "password": password}, ) if response.status_code == 401: first_run_response = await bootstrap_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 @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}")