211 lines
7.9 KiB
Python
211 lines
7.9 KiB
Python
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)
|