456 lines
14 KiB
Python
456 lines
14 KiB
Python
import asyncio
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
import json
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
from fastapi.testclient import TestClient
|
|
from openai import APIConnectionError
|
|
import pytest
|
|
|
|
from chat_api import create_app
|
|
from chat_api.config import Settings
|
|
from chat_api.service import ChatService
|
|
from chat_api.storage import JsonSessionStorage
|
|
from chat_api.storage import UserStore
|
|
|
|
|
|
@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:
|
|
def __init__(self) -> None:
|
|
self.calls: list[dict[str, Any]] = []
|
|
self.error: Exception | None = None
|
|
self.delay = 0.0
|
|
|
|
async def create(self, **kwargs: Any) -> SimpleNamespace:
|
|
self.calls.append(kwargs)
|
|
if self.delay:
|
|
await asyncio.sleep(self.delay)
|
|
if self.error:
|
|
raise self.error
|
|
return SimpleNamespace(
|
|
choices=[
|
|
SimpleNamespace(
|
|
message=SimpleNamespace(content=f"reply-{len(self.calls)}")
|
|
)
|
|
],
|
|
usage=SimpleNamespace(
|
|
prompt_tokens=10,
|
|
completion_tokens=2,
|
|
total_tokens=12,
|
|
),
|
|
)
|
|
|
|
|
|
class FakeClient:
|
|
def __init__(self) -> None:
|
|
self.completions = FakeCompletions()
|
|
self.chat = SimpleNamespace(completions=self.completions)
|
|
|
|
|
|
@pytest.fixture
|
|
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,
|
|
base_url="https://api.deepseek.com",
|
|
model="deepseek-v4-flash",
|
|
default_system_prompt="default prompt",
|
|
data_dir=tmp_path,
|
|
user_db_path=db_path,
|
|
)
|
|
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_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:
|
|
response = client.post("/sessions", json=body)
|
|
assert response.status_code == 201
|
|
return response.json()["session_id"]
|
|
|
|
|
|
def test_settings_load_api_key_from_dotenv(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> 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()
|
|
|
|
assert settings.api_key == "test-key"
|
|
|
|
|
|
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")
|
|
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_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(
|
|
app_client: AppHarness, tmp_path: Path
|
|
) -> None:
|
|
client = app_client.client
|
|
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"] == app_client.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(
|
|
app_client: AppHarness, tmp_path: Path
|
|
) -> None:
|
|
client = app_client.client
|
|
session_id = create_session(client)
|
|
|
|
data = json.loads((tmp_path / f"{session_id}.json").read_text())
|
|
assert data["user_id"] == app_client.test_user_id
|
|
assert data["system_prompt"] == "default prompt"
|
|
assert data["messages"] == []
|
|
|
|
|
|
def test_create_session_without_body(app_client: AppHarness) -> None:
|
|
client = app_client.client
|
|
|
|
response = client.post("/sessions")
|
|
|
|
assert response.status_code == 201
|
|
assert response.json()["system_prompt"] == "default prompt"
|
|
|
|
|
|
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
|
|
assert response.json()["system_prompt"] == "回答中文"
|
|
|
|
|
|
def test_multi_turn_request_contains_full_history(
|
|
app_client: AppHarness, tmp_path: Path
|
|
) -> None:
|
|
client, provider = app_client.client, app_client.provider
|
|
session_id = create_session(client, system_prompt="system")
|
|
|
|
first = client.post(
|
|
f"/sessions/{session_id}/messages", json={"content": "first"}
|
|
)
|
|
second = client.post(
|
|
f"/sessions/{session_id}/messages", json={"content": "second"}
|
|
)
|
|
|
|
assert first.status_code == second.status_code == 200
|
|
assert provider.completions.calls[1]["messages"] == [
|
|
{"role": "system", "content": "system"},
|
|
{"role": "user", "content": "first"},
|
|
{"role": "assistant", "content": "reply-1"},
|
|
{"role": "user", "content": "second"},
|
|
]
|
|
assert second.json()["usage"] == {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 2,
|
|
"total_tokens": 12,
|
|
}
|
|
assert second.json()["tools_use"] == []
|
|
datetime.fromisoformat(second.json()["message"]["created_at"])
|
|
assert "tool_calls" not in second.json()["message"]
|
|
data = json.loads((tmp_path / f"{session_id}.json").read_text())
|
|
assert [message["content"] for message in data["messages"]] == [
|
|
"first",
|
|
"reply-1",
|
|
"second",
|
|
"reply-2",
|
|
]
|
|
assert all(message["created_at"].endswith("Z") for message in data["messages"])
|
|
|
|
|
|
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",
|
|
json={"content": "hello"},
|
|
)
|
|
blank = client.post(
|
|
"/sessions/00000000-0000-0000-0000-000000000000/messages",
|
|
json={"content": " "},
|
|
)
|
|
|
|
assert missing.status_code == 404
|
|
assert blank.status_code == 422
|
|
|
|
|
|
def test_upstream_failure_does_not_change_history(
|
|
app_client: AppHarness, tmp_path: Path
|
|
) -> None:
|
|
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]
|
|
|
|
response = client.post(
|
|
f"/sessions/{session_id}/messages", json={"content": "hello"}
|
|
)
|
|
|
|
assert response.status_code == 502
|
|
assert (tmp_path / f"{session_id}.json").read_text() == before
|
|
|
|
|
|
def test_corrupt_session_returns_500(
|
|
app_client: AppHarness, tmp_path: Path
|
|
) -> None:
|
|
client = app_client.client
|
|
session_id = create_session(client)
|
|
(tmp_path / f"{session_id}.json").write_text("not-json")
|
|
|
|
response = client.post(
|
|
f"/sessions/{session_id}/messages", json={"content": "hello"}
|
|
)
|
|
|
|
assert response.status_code == 500
|
|
|
|
|
|
def test_legacy_messages_without_timestamp_remain_usable(
|
|
app_client: AppHarness, tmp_path: Path
|
|
) -> None:
|
|
client = app_client.client
|
|
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))
|
|
|
|
response = client.post(
|
|
f"/sessions/{session_id}/messages", json={"content": "new"}
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
migrated = json.loads(path.read_text())
|
|
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(app_client: AppHarness) -> None:
|
|
client = app_client.client
|
|
session_id = create_session(client)
|
|
other_headers = {"X-API-Key": app_client.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(app_client: AppHarness) -> None:
|
|
client, provider = app_client.client, app_client.provider
|
|
session_id = create_session(client)
|
|
|
|
response = client.post(
|
|
f"/sessions/{session_id}/messages",
|
|
json={"content": "hello", "stream": True},
|
|
)
|
|
|
|
assert response.status_code == 422
|
|
assert provider.completions.calls == []
|
|
|
|
|
|
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"})
|
|
|
|
response = client.get(f"/sessions/{session_id}/messages")
|
|
|
|
assert response.status_code == 200
|
|
history = response.json()
|
|
assert history["session_id"] == session_id
|
|
assert history["system_prompt"] == "history prompt"
|
|
assert history["created_at"].endswith("Z")
|
|
assert history["updated_at"].endswith("Z")
|
|
assert [message["content"] for message in history["messages"]] == [
|
|
"first",
|
|
"reply-1",
|
|
"second",
|
|
"reply-2",
|
|
]
|
|
assert all(message["created_at"].endswith("Z") for message in history["messages"])
|
|
|
|
|
|
def test_get_session_history_handles_empty_and_missing_sessions(
|
|
app_client: AppHarness
|
|
) -> None:
|
|
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")
|
|
|
|
assert empty.status_code == 200
|
|
assert empty.json()["messages"] == []
|
|
assert missing.status_code == 404
|
|
|
|
|
|
def test_same_session_concurrent_messages_are_serialized(tmp_path: Path) -> None:
|
|
async def scenario() -> None:
|
|
provider = FakeClient()
|
|
provider.completions.delay = 0.01
|
|
storage = JsonSessionStorage(tmp_path)
|
|
session = await storage.create("system")
|
|
service = ChatService(
|
|
storage=storage,
|
|
client=provider, # type: ignore[arg-type]
|
|
model="deepseek-v4-flash",
|
|
)
|
|
|
|
await asyncio.gather(
|
|
service.generate_response(session.session_id, "first"),
|
|
service.generate_response(session.session_id, "second"),
|
|
)
|
|
|
|
assert provider.completions.calls[1]["messages"] == [
|
|
{"role": "system", "content": "system"},
|
|
{"role": "user", "content": "first"},
|
|
{"role": "assistant", "content": "reply-1"},
|
|
{"role": "user", "content": "second"},
|
|
]
|
|
|
|
asyncio.run(scenario()) |