增加鉴权系统

tests/test_api.py

@@ -10,11 +10,15 @@ from openai import APIConnectionError
 import pytest
 
 from app import create_app
-from config import Settings
+from config import Settings, _parse_api_key_whitelist
 from service import ChatService
 from storage import JsonSessionStorage
 
 
+TEST_USER_ID = "11111111-1111-4111-8111-111111111111"
+OTHER_USER_ID = "22222222-2222-4222-8222-222222222222"
+
+
 class FakeCompletions:
     def __init__(self) -> None:
         self.calls: list[dict[str, Any]] = []
@@ -55,9 +59,14 @@ 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",
+        },
     )
     app = create_app(settings=settings, client=provider)  # type: ignore[arg-type]
     with TestClient(app) as client:
+        client.headers["X-API-Key"] = "test-api-key"
         yield client, provider
 
 
@@ -79,6 +88,54 @@ 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
+    del client.headers["X-API-Key"]
+
+    missing = client.post("/sessions")
+    invalid = client.get("/usage", headers={"X-API-Key": "invalid"})
+
+    assert missing.status_code == 401
+    assert missing.json() == {"detail": "api_key is required"}
+    assert invalid.status_code == 401
+    assert invalid.json() == {"detail": "api_key is required"}
+
+
+def test_usage_endpoint_returns_current_user_totals(
+    client_and_provider: tuple[TestClient, FakeClient], tmp_path: Path
+) -> None:
+    client, _ = client_and_provider
+    session_id = create_session(client)
+    message = client.post(
+        f"/sessions/{session_id}/messages", json={"content": "hello"}
+    )
+
+    response = client.get("/usage")
+
+    assert message.status_code == 200
+    assert response.status_code == 200
+    assert response.json()["user_id"] == TEST_USER_ID
+    assert response.json()["api_calls"] == 1
+    assert response.json()["prompt_tokens"] == 10
+    assert response.json()["completion_tokens"] == 2
+    assert response.json()["total_tokens"] == 12
+    assert not (tmp_path / "usage.json").exists()
+
+
 def test_create_session_uses_default_prompt(
     client_and_provider: tuple[TestClient, FakeClient], tmp_path: Path
 ) -> None:
@@ -86,6 +143,7 @@ def test_create_session_uses_default_prompt(
     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["system_prompt"] == "default prompt"
     assert data["messages"] == []
 
@@ -204,6 +262,7 @@ def test_legacy_messages_without_timestamp_remain_usable(
     session_id = create_session(client)
     path = tmp_path / f"{session_id}.json"
     data = json.loads(path.read_text())
+    data.pop("user_id")
     data["messages"] = [{"role": "user", "content": "legacy"}]
     path.write_text(json.dumps(data))
 
@@ -213,9 +272,30 @@ 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 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
+    session_id = create_session(client)
+    other_headers = {"X-API-Key": "other-api-key"}
+
+    history = client.get(
+        f"/sessions/{session_id}/messages", headers=other_headers
+    )
+    message = client.post(
+        f"/sessions/{session_id}/messages",
+        headers=other_headers,
+        json={"content": "hello"},
+    )
+
+    assert history.status_code == 404
+    assert message.status_code == 404
+
+
 def test_stream_parameter_is_rejected(
     client_and_provider: tuple[TestClient, FakeClient]
 ) -> None: