302 lines
9.3 KiB
Python
302 lines
9.3 KiB
Python
import asyncio
|
|
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 app import create_app
|
|
from config import Settings
|
|
from service import ChatService
|
|
from storage import JsonSessionStorage
|
|
|
|
|
|
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 client_and_provider(tmp_path: Path) -> tuple[TestClient, FakeClient]:
|
|
provider = FakeClient()
|
|
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,
|
|
)
|
|
app = create_app(settings=settings, client=provider) # type: ignore[arg-type]
|
|
with TestClient(app) as client:
|
|
yield client, provider
|
|
|
|
|
|
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)
|
|
(tmp_path / ".env").write_text("DEEPSEEK_API_KEY=test-key\n")
|
|
|
|
settings = Settings.from_env()
|
|
|
|
assert settings.api_key == "test-key"
|
|
|
|
|
|
def test_create_session_uses_default_prompt(
|
|
client_and_provider: tuple[TestClient, FakeClient], tmp_path: Path
|
|
) -> None:
|
|
client, _ = client_and_provider
|
|
session_id = create_session(client)
|
|
|
|
data = json.loads((tmp_path / f"{session_id}.json").read_text())
|
|
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
|
|
|
|
response = client.post("/sessions")
|
|
|
|
assert response.status_code == 201
|
|
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
|
|
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(
|
|
client_and_provider: tuple[TestClient, FakeClient], tmp_path: Path
|
|
) -> None:
|
|
client, provider = client_and_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,
|
|
}
|
|
datetime.fromisoformat(second.json()["message"]["created_at"])
|
|
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(
|
|
client_and_provider: tuple[TestClient, FakeClient]
|
|
) -> None:
|
|
client, _ = client_and_provider
|
|
|
|
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(
|
|
client_and_provider: tuple[TestClient, FakeClient], tmp_path: Path
|
|
) -> None:
|
|
client, provider = client_and_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(
|
|
client_and_provider: tuple[TestClient, FakeClient], tmp_path: Path
|
|
) -> None:
|
|
client, _ = client_and_provider
|
|
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(
|
|
client_and_provider: tuple[TestClient, FakeClient], tmp_path: Path
|
|
) -> None:
|
|
client, _ = client_and_provider
|
|
session_id = create_session(client)
|
|
path = tmp_path / f"{session_id}.json"
|
|
data = json.loads(path.read_text())
|
|
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 all("created_at" in message for message in migrated["messages"])
|
|
|
|
|
|
def test_stream_parameter_is_rejected(
|
|
client_and_provider: tuple[TestClient, FakeClient]
|
|
) -> None:
|
|
client, provider = client_and_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(
|
|
client_and_provider: tuple[TestClient, FakeClient]
|
|
) -> None:
|
|
client, _ = client_and_provider
|
|
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(
|
|
client_and_provider: tuple[TestClient, FakeClient]
|
|
) -> None:
|
|
client, _ = client_and_provider
|
|
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())
|