Serendipity
d021390fd3
- Implement FastAPI server with health, collections, ingest, search, and delete endpoints - Add Pydantic request models for ingest and search - Add lazy singleton initialization for DB, embedder, and services - Support environment variable override for data dir and collection name - Add integration tests with TestClient and temp directory - Set TRANSFORMERS_OFFLINE/HF_HUB_OFFLINE for offline model loading in tests
93 lines
2.6 KiB
Python
93 lines
2.6 KiB
Python
"""API 端点集成测试."""
|
|
import os
|
|
import tempfile
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
|
|
@pytest.fixture
|
|
def client():
|
|
"""创建测试客户端（使用临时目录 + 本地嵌入模型）."""
|
|
tmpdir = tempfile.mkdtemp()
|
|
# 设置环境变量使 app 使用测试配置
|
|
os.environ["MD_VECTOR_DB_DATA_DIR"] = tmpdir
|
|
os.environ["MD_VECTOR_DB_COLLECTION"] = "test_api"
|
|
# 离线模式: 模型已缓存, 无需联网
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
|
|
|
from src.server.app import app
|
|
|
|
with TestClient(app) as c:
|
|
yield c
|
|
|
|
|
|
class TestHealthEndpoint:
|
|
"""健康检查."""
|
|
|
|
def test_health_returns_ok(self, client):
|
|
response = client.get("/api/v1/health")
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["status"] == "ok"
|
|
|
|
|
|
class TestCollectionEndpoint:
|
|
"""Collection 端点."""
|
|
|
|
def test_list_collections(self, client):
|
|
response = client.get("/api/v1/collections")
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "collections" in data
|
|
|
|
|
|
class TestSearchEndpoint:
|
|
"""搜索端点."""
|
|
|
|
def test_search_requires_query(self, client):
|
|
response = client.post("/api/v1/search", json={})
|
|
assert response.status_code == 422 # FastAPI 自动校验
|
|
|
|
def test_search_empty_collection(self, client):
|
|
response = client.post(
|
|
"/api/v1/search",
|
|
json={"query": "test", "top_k": 5},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "results" in data
|
|
assert data["results"] == []
|
|
|
|
|
|
class TestIngestEndpoint:
|
|
"""入库端点."""
|
|
|
|
def test_ingest_content(self, client):
|
|
response = client.post(
|
|
"/api/v1/ingest",
|
|
json={
|
|
"content": "# Test\nHello world.",
|
|
"file_name": "test.md",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["status"] == "ok"
|
|
assert data["chunks"] > 0
|
|
|
|
def test_ingest_then_search(self, client):
|
|
"""入库后能检索到."""
|
|
client.post(
|
|
"/api/v1/ingest",
|
|
json={"content": "# 配置说明\nChromaDB 配置很简单。", "file_name": "config.md"},
|
|
)
|
|
response = client.post(
|
|
"/api/v1/search",
|
|
json={"query": "配置", "top_k": 3},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert len(data["results"]) > 0
|