141 lines
4.2 KiB
Python
141 lines
4.2 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
|
|
|
|
|
|
class TestSecurity:
|
|
"""安全测试."""
|
|
|
|
def test_ingest_rejects_path_traversal(self, client):
|
|
"""拒绝路径遍历攻击."""
|
|
response = client.post(
|
|
"/api/v1/ingest",
|
|
json={"file_path": "../../../etc/passwd"},
|
|
)
|
|
assert response.status_code in (400, 403)
|
|
|
|
def test_ingest_rejects_absolute_path(self, client):
|
|
"""拒绝绝对路径."""
|
|
response = client.post(
|
|
"/api/v1/ingest",
|
|
json={"file_path": "C:\\Windows\\System32\\config\\SAM"},
|
|
)
|
|
assert response.status_code in (400, 403)
|
|
|
|
|
|
class TestDeleteEndpoint:
|
|
"""删除端点."""
|
|
|
|
def test_delete_nonexistent(self, client):
|
|
"""删除不存在的文件返回 404."""
|
|
response = client.delete("/api/v1/documents/nonexistent.md")
|
|
assert response.status_code == 404
|
|
|
|
def test_delete_ingested(self, client):
|
|
"""删除已入库文件后搜索不再返回."""
|
|
# 入库
|
|
ingest_resp = client.post(
|
|
"/api/v1/ingest",
|
|
json={"content": "# Test Delete\nHello.", "file_name": "delete-test.md"},
|
|
)
|
|
assert ingest_resp.status_code == 200, f"ingest failed: {ingest_resp.json()}"
|
|
# 删除
|
|
response = client.delete("/api/v1/documents/delete-test.md")
|
|
assert response.status_code == 200, f"delete failed: {response.json()}"
|
|
# 搜索验证已删除
|
|
search_resp = client.post(
|
|
"/api/v1/search", json={"query": "Test Delete", "top_k": 3}
|
|
)
|
|
results = search_resp.json()["results"]
|
|
sources = [r["source_file"] for r in results]
|
|
assert "delete-test.md" not in sources
|