""" Shared pytest fixtures for exercising FastAPI routers over the running API service. """ from __future__ import annotations import os import uuid from collections.abc import AsyncGenerator import anyio import httpx import pytest import pytest_asyncio from dotenv import load_dotenv # Load project and test specific environment variables. load_dotenv(".env", override=False) load_dotenv("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") assert ADMIN_LOGIN, "TEST_USERNAME is not set" assert ADMIN_PASSWORD, "TEST_PASSWORD is not set" _ADMIN_TOKEN_CACHE: str | None = None HTTP_TIMEOUT = httpx.Timeout(30.0, connect=5.0) @pytest_asyncio.fixture(scope="function") async def test_client() -> AsyncGenerator[httpx.AsyncClient, None]: """Async HTTP client bound to the live API base URL.""" 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: """Authenticate with super admin credentials and cache the bearer token.""" global _ADMIN_TOKEN_CACHE if _ADMIN_TOKEN_CACHE: return _ADMIN_TOKEN_CACHE if not ADMIN_LOGIN or not ADMIN_PASSWORD: pytest.skip("Admin credentials are not configured via environment variables.") 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 == 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. Please complete `/api/auth/initialize` " "before running router 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]: """Authorization headers for the super admin user.""" return {"Authorization": f"Bearer {admin_token}"} @pytest.fixture(scope="session", autouse=True) def cleanup_test_knowledge_databases(): """ Best-effort cleanup for leftover test knowledge databases (e.g. pytest_* / py_test*). """ 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 e: print(f"Warning: Failed to list knowledge databases for cleanup: {e}") 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 "" db_id = entry.get("db_id") if not db_id or not isinstance(name, str) or not name.startswith(prefixes): continue try: delete_response = await client.delete(f"/api/knowledge/databases/{db_id}", headers=headers) if delete_response.status_code not in (200, 404): print(f"Warning: Failed to cleanup knowledge database {db_id}: {delete_response.text}") except Exception as e: print(f"Warning: Exception during cleanup of {db_id}: {e}") try: anyio.run(run_cleanup) except Exception as e: print(f"Warning: Exception during session cleanup startup: {e}") yield try: anyio.run(run_cleanup) except Exception as e: print(f"Warning: Exception during session cleanup teardown: {e}") @pytest_asyncio.fixture(scope="function") async def standard_user(test_client: httpx.AsyncClient, admin_headers: dict[str, str]) -> dict: """ Provision a temporary standard user for permission checks. Yields a dictionary with `user`, `password`, and `headers` keys. """ username = f"pytest_user_{uuid.uuid4().hex[:8]}" password = f"Pw!{uuid.uuid4().hex[:8]}" response = await test_client.post( "/api/auth/users", json={"username": username, "password": password, "role": "user"}, 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["user_id"], "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: response = await test_client.delete(f"/api/auth/users/{user_payload['id']}", headers=admin_headers) assert response.status_code == 200, f"Failed to cleanup test user {user_payload['user_id']}: {response.text}" @pytest_asyncio.fixture(scope="function") async def knowledge_database(test_client: httpx.AsyncClient, admin_headers: dict[str, str]) -> dict: """ Create a temporary knowledge database for tests that need LightRAG metadata. """ import time # 使用UUID作为数据库名称的一部分，确保唯一性 unique_id = uuid.uuid4().hex timestamp = int(time.time() * 1000000) # 微秒级时间戳 db_name = f"pytest_kb_{timestamp}_{unique_id}" db_id = None try: create_response = await test_client.post( "/api/knowledge/databases", json={ "database_name": db_name, "description": "Pytest managed knowledge base", "embed_model_name": "siliconflow/BAAI/bge-m3", "kb_type": "lightrag", "additional_params": {}, }, headers=admin_headers, ) if create_response.status_code == 200: db_payload = create_response.json() db_id = db_payload["db_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 if db_id else {"db_id": db_id, "name": db_name} finally: # 确保清理，即使测试失败 if db_id: try: delete_response = await test_client.delete(f"/api/knowledge/databases/{db_id}", headers=admin_headers) if delete_response.status_code != 200: print(f"Warning: Failed to cleanup knowledge database {db_id}: {delete_response.text}") except Exception as e: print(f"Warning: Exception during cleanup of {db_id}: {e}") def pytest_configure(config: pytest.Config) -> None: """Register commonly used custom markers.""" config.addinivalue_line("markers", "auth: marks tests that require authentication") config.addinivalue_line("markers", "integration: marks tests that hit the live API service") pytest_plugins = ["pytest_asyncio"]