Mplan/simple-chat-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

优化鉴权系统

Browse Source
This commit is contained in:
Mplan
2026-07-03 21:54:19 +08:00
parent 7e65fadd83
commit 98a92b83fa
11 changed files with 539 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_api.py
+152 -79
View File
@@ -1,4 +1,5 @@
import asyncio
from dataclasses import dataclass
from datetime import datetime
import json
from pathlib import Path
@@ -10,13 +11,20 @@ from openai import APIConnectionError
import pytest
from app import create_app
from config import Settings, _parse_api_key_whitelist
from config import Settings
from service import ChatService
from storage import JsonSessionStorage
from users import UserStore
TEST_USER_ID = "11111111-1111-4111-8111-111111111111"
OTHER_USER_ID = "22222222-2222-4222-8222-222222222222"
@dataclass
class AppHarness:
client: TestClient
provider: "FakeClient"
test_user_id: str
other_user_id: str
test_api_key: str
other_api_key: str
class FakeCompletions:
@@ -43,6 +51,8 @@ class FakeCompletions:
total_tokens=12,
),
)
class FakeClient:
def __init__(self) -> None:
self.completions = FakeCompletions()
@@ -50,8 +60,14 @@ class FakeClient:
@pytest.fixture
def client_and_provider(tmp_path: Path) -> tuple[TestClient, FakeClient]:
def app_client(tmp_path: Path) -> AppHarness:
provider = FakeClient()
db_path = tmp_path / "users.db"
store = UserStore(db_path)
asyncio.run(store.init())
test_user = asyncio.run(store.register("test-user", "test-password"))
other_user = asyncio.run(store.register("other-user", "other-password"))
settings = Settings(
port=8000,
api_key=None,
@@ -59,15 +75,19 @@ def client_and_provider(tmp_path: Path) -> tuple[TestClient, FakeClient]:
model="deepseek-v4-flash",
default_system_prompt="default prompt",
data_dir=tmp_path,
api_key_whitelist={
TEST_USER_ID: "test-api-key",
OTHER_USER_ID: "other-api-key",
},
user_db_path=db_path,
)
app = create_app(settings=settings, client=provider) # type: ignore[arg-type]
app = create_app(settings=settings, client=provider, user_store=store) # type: ignore[arg-type]
with TestClient(app) as client:
client.headers["X-API-Key"] = "test-api-key"
yield client, provider
client.headers["X-API-Key"] = test_user.api_key
yield AppHarness(
client=client,
provider=provider,
test_user_id=test_user.user_id,
other_user_id=other_user.user_id,
test_api_key=test_user.api_key,
other_api_key=other_user.api_key,
)
def create_session(client: TestClient, **body: str) -> str:
@@ -81,6 +101,8 @@ def test_settings_load_api_key_from_dotenv(
) -> None:
monkeypatch.chdir(tmp_path)
monkeypatch.delenv("DEEPSEEK_API_KEY", raising=False)
monkeypatch.delenv("USER_DB_PATH", raising=False)
monkeypatch.delenv("CHAT_DATA_DIR", raising=False)
(tmp_path / ".env").write_text("DEEPSEEK_API_KEY=test-key\n")
settings = Settings.from_env()
@@ -88,22 +110,8 @@ def test_settings_load_api_key_from_dotenv(
assert settings.api_key == "test-key"
def test_api_key_whitelist_requires_and_normalizes_uuid() -> None:
parsed = _parse_api_key_whitelist(
'{"7A007F92-86D0-4F51-931F-007B9F0B8B0E":"key"}'
)
assert parsed == {
"7a007f92-86d0-4f51-931f-007b9f0b8b0e": "key"
}
with pytest.raises(ValueError, match="valid UUIDs"):
_parse_api_key_whitelist('{"custom-name":"key"}')
def test_all_api_endpoints_require_valid_api_key(
client_and_provider: tuple[TestClient, FakeClient]
) -> None:
client, _ = client_and_provider
def test_all_api_endpoints_require_valid_api_key(app_client: AppHarness) -> None:
client = app_client.client
del client.headers["X-API-Key"]
missing = client.post("/sessions")
@@ -115,10 +123,93 @@ def test_all_api_endpoints_require_valid_api_key(
assert invalid.json() == {"detail": "api_key is required"}
def test_register_creates_user_and_returns_user_id(app_client: AppHarness) -> None:
client = app_client.client
response = client.post(
"/auth/register",
json={"username": "fresh-user", "password": "fresh-password"},
)
assert response.status_code == 201
body = response.json()
assert isinstance(body["user_id"], str)
assert body["user_id"] != app_client.test_user_id
assert body["created_at"].endswith("Z")
datetime.fromisoformat(body["created_at"])
login = client.post(
"/auth/login",
json={"username": "fresh-user", "password": "fresh-password"},
)
assert login.status_code == 200
assert login.json()["user_id"] == body["user_id"]
assert isinstance(login.json()["api_key"], str)
assert login.json()["api_key"]
def test_register_rejects_duplicate_username_with_409(app_client: AppHarness) -> None:
client = app_client.client
response = client.post(
"/auth/register",
json={"username": "test-user", "password": "any-password"},
)
assert response.status_code == 409
assert response.json() == {"detail": "username already exists"}
def test_register_rejects_blank_fields(app_client: AppHarness) -> None:
client = app_client.client
blank_username = client.post(
"/auth/register", json={"username": " ", "password": "secret"}
)
blank_password = client.post(
"/auth/register", json={"username": "new-user", "password": " "}
)
assert blank_username.status_code == 422
assert blank_password.status_code == 422
def test_login_returns_persistent_api_key(app_client: AppHarness) -> None:
client = app_client.client
response = client.post(
"/auth/login",
json={"username": "test-user", "password": "test-password"},
)
assert response.status_code == 200
body = response.json()
assert body["user_id"] == app_client.test_user_id
assert body["api_key"] == app_client.test_api_key
def test_login_invalid_credentials_returns_401(app_client: AppHarness) -> None:
client = app_client.client
wrong_password = client.post(
"/auth/login",
json={"username": "test-user", "password": "wrong"},
)
unknown_user = client.post(
"/auth/login",
json={"username": "no-such-user", "password": "anything"},
)
assert wrong_password.status_code == 401
assert wrong_password.json() == {"detail": "invalid username or password"}
assert unknown_user.status_code == 401
assert unknown_user.json() == {"detail": "invalid username or password"}
def test_usage_endpoint_returns_current_user_totals(
client_and_provider: tuple[TestClient, FakeClient], tmp_path: Path
app_client: AppHarness, tmp_path: Path
) -> None:
client, _ = client_and_provider
client = app_client.client
session_id = create_session(client)
message = client.post(
f"/sessions/{session_id}/messages", json={"content": "hello"}
@@ -128,7 +219,7 @@ def test_usage_endpoint_returns_current_user_totals(
assert message.status_code == 200
assert response.status_code == 200
assert response.json()["user_id"] == TEST_USER_ID
assert response.json()["user_id"] == app_client.test_user_id
assert response.json()["api_calls"] == 1
assert response.json()["prompt_tokens"] == 10
assert response.json()["completion_tokens"] == 2
@@ -137,21 +228,19 @@ def test_usage_endpoint_returns_current_user_totals(
def test_create_session_uses_default_prompt(
client_and_provider: tuple[TestClient, FakeClient], tmp_path: Path
app_client: AppHarness, tmp_path: Path
) -> None:
client, _ = client_and_provider
client = app_client.client
session_id = create_session(client)
data = json.loads((tmp_path / f"{session_id}.json").read_text())
assert data["user_id"] == TEST_USER_ID
assert data["user_id"] == app_client.test_user_id
assert data["system_prompt"] == "default prompt"
assert data["messages"] == []
def test_create_session_without_body(
client_and_provider: tuple[TestClient, FakeClient]
) -> None:
client, _ = client_and_provider
def test_create_session_without_body(app_client: AppHarness) -> None:
client = app_client.client
response = client.post("/sessions")
@@ -159,10 +248,8 @@ def test_create_session_without_body(
assert response.json()["system_prompt"] == "default prompt"
def test_create_session_accepts_custom_prompt(
client_and_provider: tuple[TestClient, FakeClient]
) -> None:
client, _ = client_and_provider
def test_create_session_accepts_custom_prompt(app_client: AppHarness) -> None:
client = app_client.client
response = client.post("/sessions", json={"system_prompt": "回答中文"})
assert response.status_code == 201
@@ -170,9 +257,9 @@ def test_create_session_accepts_custom_prompt(
def test_multi_turn_request_contains_full_history(
client_and_provider: tuple[TestClient, FakeClient], tmp_path: Path
app_client: AppHarness, tmp_path: Path
) -> None:
client, provider = client_and_provider
client, provider = app_client.client, app_client.provider
session_id = create_session(client, system_prompt="system")
first = client.post(
@@ -207,10 +294,8 @@ def test_multi_turn_request_contains_full_history(
assert all(message["created_at"].endswith("Z") for message in data["messages"])
def test_invalid_session_and_blank_message(
client_and_provider: tuple[TestClient, FakeClient]
) -> None:
client, _ = client_and_provider
def test_invalid_session_and_blank_message(app_client: AppHarness) -> None:
client = app_client.client
missing = client.post(
"/sessions/00000000-0000-0000-0000-000000000000/messages",
@@ -226,9 +311,9 @@ def test_invalid_session_and_blank_message(
def test_upstream_failure_does_not_change_history(
client_and_provider: tuple[TestClient, FakeClient], tmp_path: Path
app_client: AppHarness, tmp_path: Path
) -> None:
client, provider = client_and_provider
client, provider = app_client.client, app_client.provider
session_id = create_session(client)
before = (tmp_path / f"{session_id}.json").read_text()
provider.completions.error = APIConnectionError(request=object()) # type: ignore[arg-type]
@@ -242,9 +327,9 @@ def test_upstream_failure_does_not_change_history(
def test_corrupt_session_returns_500(
client_and_provider: tuple[TestClient, FakeClient], tmp_path: Path
app_client: AppHarness, tmp_path: Path
) -> None:
client, _ = client_and_provider
client = app_client.client
session_id = create_session(client)
(tmp_path / f"{session_id}.json").write_text("not-json")
@@ -256,9 +341,9 @@ def test_corrupt_session_returns_500(
def test_legacy_messages_without_timestamp_remain_usable(
client_and_provider: tuple[TestClient, FakeClient], tmp_path: Path
app_client: AppHarness, tmp_path: Path
) -> None:
client, _ = client_and_provider
client = app_client.client
session_id = create_session(client)
path = tmp_path / f"{session_id}.json"
data = json.loads(path.read_text())
@@ -272,16 +357,14 @@ def test_legacy_messages_without_timestamp_remain_usable(
assert response.status_code == 200
migrated = json.loads(path.read_text())
assert migrated["user_id"] == TEST_USER_ID
assert migrated["user_id"] == app_client.test_user_id
assert all("created_at" in message for message in migrated["messages"])
def test_session_cannot_be_accessed_by_another_user(
client_and_provider: tuple[TestClient, FakeClient]
) -> None:
client, _ = client_and_provider
def test_session_cannot_be_accessed_by_another_user(app_client: AppHarness) -> None:
client = app_client.client
session_id = create_session(client)
other_headers = {"X-API-Key": "other-api-key"}
other_headers = {"X-API-Key": app_client.other_api_key}
history = client.get(
f"/sessions/{session_id}/messages", headers=other_headers
@@ -296,10 +379,8 @@ def test_session_cannot_be_accessed_by_another_user(
assert message.status_code == 404
def test_stream_parameter_is_rejected(
client_and_provider: tuple[TestClient, FakeClient]
) -> None:
client, provider = client_and_provider
def test_stream_parameter_is_rejected(app_client: AppHarness) -> None:
client, provider = app_client.client, app_client.provider
session_id = create_session(client)
response = client.post(
@@ -311,17 +392,11 @@ def test_stream_parameter_is_rejected(
assert provider.completions.calls == []
def test_get_session_history_returns_all_messages(
client_and_provider: tuple[TestClient, FakeClient]
) -> None:
client, _ = client_and_provider
def test_get_session_history_returns_all_messages(app_client: AppHarness) -> None:
client = app_client.client
session_id = create_session(client, system_prompt="history prompt")
client.post(
f"/sessions/{session_id}/messages", json={"content": "first"}
)
client.post(
f"/sessions/{session_id}/messages", json={"content": "second"}
)
client.post(f"/sessions/{session_id}/messages", json={"content": "first"})
client.post(f"/sessions/{session_id}/messages", json={"content": "second"})
response = client.get(f"/sessions/{session_id}/messages")
@@ -341,15 +416,13 @@ def test_get_session_history_returns_all_messages(
def test_get_session_history_handles_empty_and_missing_sessions(
client_and_provider: tuple[TestClient, FakeClient]
app_client: AppHarness
) -> None:
client, _ = client_and_provider
client = app_client.client
session_id = create_session(client)
empty = client.get(f"/sessions/{session_id}/messages")
missing = client.get(
"/sessions/00000000-0000-0000-0000-000000000000/messages"
)
missing = client.get("/sessions/00000000-0000-0000-0000-000000000000/messages")
assert empty.status_code == 200
assert empty.json()["messages"] == []
@@ -380,4 +453,4 @@ def test_same_session_concurrent_messages_are_serialized(tmp_path: Path) -> None
{"role": "user", "content": "second"},
]
asyncio.run(scenario())
asyncio.run(scenario())