from __future__ import annotations import sqlite3 import httpx from backend.bootstrap.application import create_application from backend.bootstrap.settings import Settings from backend.llm.provider import ProviderFailure from tests.support import run_scenario from tests.test_accounts import ( ADMIN_PASSWORD, USER_PASSWORD, csrf_headers, register, use_session, ) def model_payload(index: int, api_key: str | None = None) -> dict[str, str]: return { "display_name": f"模型 {index}", "base_url": f"https://model-{index}.example.com/v1", "model_identifier": f"model-{index}", "api_key": api_key or f"secret-key-{index}", } class ScriptedProvider: def __init__(self, scripts: list[object]) -> None: self.scripts = scripts def stream(self, _profile, _messages): script = self.scripts.pop(0) if isinstance(script, Exception): raise script yield from script def test_model_pool_is_admin_only_and_never_exposes_keys(tmp_path) -> None: application = create_application(Settings.for_test(tmp_path)) async def scenario(client: httpx.AsyncClient) -> None: _, admin_session = await register(client, "model-admin", ADMIN_PASSWORD) client.cookies.clear() _, user_session = await register(client, "model-user", USER_PASSWORD) denied_read = await client.get("/api/admin/models") assert denied_read.status_code == 403 denied_write = await client.post( "/api/admin/models", headers=csrf_headers(user_session), json=model_payload(1), ) assert denied_write.status_code == 403 use_session(client, admin_session) created = await client.post( "/api/admin/models", headers=csrf_headers(admin_session), json=model_payload(1, "private-alpha-key"), ) assert created.status_code == 201 assert created.json()["is_primary"] is True assert created.json()["is_fallback"] is False assert created.json()["has_api_key"] is True assert "api_key" not in created.json() assert "private-alpha-key" not in created.text with sqlite3.connect(application.state.settings.database_path) as connection: encrypted_before = connection.execute( "SELECT encrypted_api_key FROM llm_models WHERE id = 1" ).fetchone()[0] assert encrypted_before != "private-alpha-key" assert "private-alpha-key" not in encrypted_before application.state.container.llm._provider = ScriptedProvider( [["连接", "成功"], ProviderFailure("authentication")] ) tested = await client.post( "/api/admin/models/1/test", headers=csrf_headers(admin_session) ) assert tested.status_code == 200 assert tested.json()["connected"] is True assert tested.json()["message"] == "连接成功" assert "private-alpha-key" not in tested.text failed_test = await client.post( "/api/admin/models/1/test", headers=csrf_headers(admin_session) ) assert failed_test.status_code == 503 assert failed_test.json()["error"]["code"] == "model_authentication" with sqlite3.connect(application.state.settings.database_path) as connection: requests = connection.execute( """ SELECT feature, business_id, status, error_type FROM llm_requests ORDER BY started_at, rowid """ ).fetchall() attempts = connection.execute( """ SELECT model_id, role, status, error_type FROM llm_attempts ORDER BY id """ ).fetchall() usage = connection.execute("SELECT COUNT(*) FROM llm_usage_daily").fetchone()[0] assert requests == [ ("model_connectivity", "model:1", "success", ""), ("model_connectivity", "model:1", "failed", "authentication"), ] assert attempts == [ (1, "primary", "success", ""), (1, "primary", "failed", "authentication"), ] assert usage == 0 updated_payload = model_payload(1) updated_payload.pop("api_key") updated_payload["display_name"] = "主模型" updated = await client.put( "/api/admin/models/1", headers=csrf_headers(admin_session), json=updated_payload, ) assert updated.status_code == 200 assert updated.json()["display_name"] == "主模型" with sqlite3.connect(application.state.settings.database_path) as connection: encrypted_after = connection.execute( "SELECT encrypted_api_key FROM llm_models WHERE id = 1" ).fetchone()[0] assert encrypted_after == encrypted_before run_scenario(application, scenario) def test_model_selection_deletion_guards_and_runtime_config(tmp_path) -> None: application = create_application(Settings.for_test(tmp_path)) async def scenario(client: httpx.AsyncClient) -> None: _, admin_session = await register(client, "selection-admin", ADMIN_PASSWORD) for index in (1, 2, 3): response = await client.post( "/api/admin/models", headers=csrf_headers(admin_session), json=model_payload(index), ) assert response.status_code == 201 duplicate_roles = await client.put( "/api/admin/models/selection", headers=csrf_headers(admin_session), json={"primary_model_id": 2, "fallback_model_id": 2}, ) assert duplicate_roles.status_code == 400 assert duplicate_roles.json()["error"]["code"] == "duplicate_model_role" selected = await client.put( "/api/admin/models/selection", headers=csrf_headers(admin_session), json={"primary_model_id": 2, "fallback_model_id": 1}, ) assert selected.status_code == 200 models = (await client.get("/api/admin/models")).json() assert next(item for item in models if item["id"] == 2)["is_primary"] is True assert next(item for item in models if item["id"] == 1)["is_fallback"] is True selected_delete = await client.delete( "/api/admin/models/1", headers=csrf_headers(admin_session) ) assert selected_delete.status_code == 409 assert selected_delete.json()["error"]["code"] == "model_in_use" unselected_delete = await client.delete( "/api/admin/models/3", headers=csrf_headers(admin_session) ) assert unselected_delete.status_code == 200 runtime = application.state.container.model_pool.runtime_config() assert runtime.primary.id == 2 assert runtime.fallback is not None assert runtime.fallback.id == 1 assert ( application.state.container.model_pool.decrypt_api_key(runtime.primary) == "secret-key-2" ) run_scenario(application, scenario) def test_model_pool_rejects_twenty_first_model(tmp_path) -> None: application = create_application(Settings.for_test(tmp_path)) async def scenario(client: httpx.AsyncClient) -> None: _, admin_session = await register(client, "capacity-admin", ADMIN_PASSWORD) for index in range(1, 21): response = await client.post( "/api/admin/models", headers=csrf_headers(admin_session), json=model_payload(index), ) assert response.status_code == 201 rejected = await client.post( "/api/admin/models", headers=csrf_headers(admin_session), json=model_payload(21), ) assert rejected.status_code == 409 assert rejected.json()["error"]["code"] == "model_pool_full" run_scenario(application, scenario)