test: 添加测试框架和全面的单元测试
- 添加 pytest 配置和测试依赖到 requirements.txt - 创建测试包结构和 fixtures (conftest.py) - 添加数据库模块的 CRUD 操作测试 (test_database.py) - 添加 Provider 插件系统测试 (test_providers.py) - 添加调度器模块测试 (test_scheduler.py) - 添加 API 路由测试 (test_api.py) - 添加回归测试覆盖边界条件和错误处理 (test_regressions.py) - 添加健康检查端点用于容器监控 - 修复调度器中的日历计算逻辑和任务执行参数处理 - 更新数据库函数以返回操作结果状态
This commit is contained in:
congsh
@@ -0,0 +1 @@
|
||||
# Tests package
|
||||
@@ -0,0 +1,77 @@
|
||||
"""测试配置和 fixtures"""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from app.config import AppConfig, load_config
|
||||
from app.database import close_db, get_db
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def event_loop():
|
||||
"""创建事件循环"""
|
||||
loop = asyncio.get_event_loop_policy().new_event_loop()
|
||||
yield loop
|
||||
loop.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def temp_db():
|
||||
"""临时数据库 fixture"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 使用唯一的文件名确保每个测试都使用新数据库
|
||||
db_path = Path(tmpdir) / f"test_{uuid.uuid4().hex}.db"
|
||||
# 修改配置使用临时数据库
|
||||
import app.config as config_module
|
||||
original_db_path = config_module.settings.database.path
|
||||
config_module.settings.database.path = str(db_path)
|
||||
|
||||
# 确保使用新的数据库连接
|
||||
import app.database as db_module
|
||||
db_module._db = None
|
||||
|
||||
yield db_path
|
||||
|
||||
# 清理
|
||||
await close_db()
|
||||
config_module.settings.database.path = original_db_path
|
||||
db_module._db = None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def temp_storage():
|
||||
"""临时存储目录 fixture"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
import app.config as config_module
|
||||
original_storage_path = config_module.settings.storage.path
|
||||
config_module.settings.storage.path = tmpdir
|
||||
|
||||
yield tmpdir
|
||||
|
||||
config_module.settings.storage.path = original_storage_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def db(temp_db):
|
||||
"""初始化数据库 fixture"""
|
||||
from app import database as db_module
|
||||
await get_db()
|
||||
return db_module
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_plan():
|
||||
"""示例 Plan 数据"""
|
||||
return {
|
||||
"name": "Test Plan",
|
||||
"provider_name": "openai",
|
||||
"api_key": "sk-test-key",
|
||||
"api_base": "https://api.openai.com/v1",
|
||||
"plan_type": "coding",
|
||||
"supported_models": ["gpt-4", "gpt-3.5-turbo"],
|
||||
"enabled": True,
|
||||
}
|
||||
@@ -0,0 +1,412 @@
|
||||
"""API 路由测试"""
|
||||
|
||||
import pytest
|
||||
from httpx import AsyncClient, ASGITransport
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from app.main import app
|
||||
|
||||
|
||||
class TestHealthCheck:
|
||||
"""健康检查测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_endpoint(self):
|
||||
"""测试 /health 端点"""
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.get("/health")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "healthy"
|
||||
assert "service" in data
|
||||
|
||||
|
||||
class TestPlansAPI:
|
||||
"""Plans API 测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_plans_empty(self, temp_db, temp_storage):
|
||||
"""测试列出空的 Plans"""
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.get("/api/plans")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_plan(self, temp_db, temp_storage):
|
||||
"""测试创建 Plan"""
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post(
|
||||
"/api/plans",
|
||||
json={
|
||||
"name": "Test Plan",
|
||||
"provider_name": "openai",
|
||||
"api_key": "sk-test",
|
||||
"api_base": "https://api.openai.com/v1",
|
||||
"plan_type": "coding",
|
||||
"supported_models": ["gpt-4"],
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "id" in data
|
||||
assert data["name"] == "Test Plan"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_plan_invalid(self, temp_db, temp_storage):
|
||||
"""测试创建 Plan 无效数据"""
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post("/api/plans", json={})
|
||||
|
||||
# 应该返回 422 验证错误
|
||||
assert response.status_code == 422
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_plan(self, temp_db, temp_storage):
|
||||
"""测试获取 Plan"""
|
||||
# 先创建
|
||||
from app import database as db
|
||||
plan = await db.create_plan(
|
||||
name="Test Plan",
|
||||
provider_name="openai",
|
||||
api_key="sk-test",
|
||||
)
|
||||
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.get(f"/api/plans/{plan['id']}")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == plan["id"]
|
||||
assert data["name"] == "Test Plan"
|
||||
assert "quota_rules" in data
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_plan_not_found(self, temp_db, temp_storage):
|
||||
"""测试获取不存在的 Plan"""
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.get("/api/plans/nonexistent")
|
||||
|
||||
assert response.status_code == 404
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_plan(self, temp_db, temp_storage):
|
||||
"""测试更新 Plan"""
|
||||
from app import database as db
|
||||
plan = await db.create_plan(
|
||||
name="Old Name",
|
||||
provider_name="openai",
|
||||
api_key="sk-test",
|
||||
)
|
||||
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.patch(
|
||||
f"/api/plans/{plan['id']}",
|
||||
json={"name": "New Name"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["ok"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_plan(self, temp_db, temp_storage):
|
||||
"""测试删除 Plan"""
|
||||
from app import database as db
|
||||
plan = await db.create_plan(
|
||||
name="To Delete",
|
||||
provider_name="openai",
|
||||
api_key="sk-test",
|
||||
)
|
||||
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.delete(f"/api/plans/{plan['id']}")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["ok"] is True
|
||||
|
||||
|
||||
class TestQuotaRulesAPI:
|
||||
"""QuotaRules API 测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_quota_rule(self, temp_db, temp_storage):
|
||||
"""测试创建 QuotaRule"""
|
||||
from app import database as db
|
||||
plan = await db.create_plan(name="Test", provider_name="openai", api_key="sk-test")
|
||||
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post(
|
||||
f"/api/plans/{plan['id']}/rules",
|
||||
json={
|
||||
"rule_name": "Daily Limit",
|
||||
"quota_total": 100,
|
||||
"quota_unit": "requests",
|
||||
"refresh_type": "calendar_cycle",
|
||||
"calendar_unit": "daily",
|
||||
"calendar_anchor": {"hour": 0},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "id" in data
|
||||
assert data["rule_name"] == "Daily Limit"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_quota_rules(self, temp_db, temp_storage):
|
||||
"""测试列出 QuotaRules"""
|
||||
from app import database as db
|
||||
plan = await db.create_plan(name="Test", provider_name="openai", api_key="sk-test")
|
||||
await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule 1", quota_total=100)
|
||||
await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule 2", quota_total=200)
|
||||
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.get(f"/api/plans/{plan['id']}/rules")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_quota_rule(self, temp_db, temp_storage):
|
||||
"""测试更新 QuotaRule"""
|
||||
from app import database as db
|
||||
plan = await db.create_plan(name="Test", provider_name="openai", api_key="sk-test")
|
||||
rule = await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule", quota_total=100)
|
||||
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.patch(
|
||||
f"/api/plans/rules/{rule['id']}",
|
||||
json={"quota_total": 200},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["ok"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_quota_rule(self, temp_db, temp_storage):
|
||||
"""测试删除 QuotaRule"""
|
||||
from app import database as db
|
||||
plan = await db.create_plan(name="Test", provider_name="openai", api_key="sk-test")
|
||||
rule = await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule", quota_total=100)
|
||||
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.delete(f"/api/plans/rules/{rule['id']}")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["ok"] is True
|
||||
|
||||
|
||||
class TestModelRoutesAPI:
|
||||
"""ModelRoutes API 测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_model_route(self, temp_db, temp_storage):
|
||||
"""测试设置模型路由"""
|
||||
from app import database as db
|
||||
plan = await db.create_plan(name="Test", provider_name="openai", api_key="sk-test")
|
||||
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post(
|
||||
"/api/plans/routes/models",
|
||||
json={"model_name": "gpt-4", "plan_id": plan["id"], "priority": 10},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "id" in data
|
||||
assert data["model_name"] == "gpt-4"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_model_routes(self, temp_db, temp_storage):
|
||||
"""测试列出模型路由"""
|
||||
from app import database as db
|
||||
plan = await db.create_plan(name="Test", provider_name="openai", api_key="sk-test")
|
||||
await db.set_model_route("gpt-4", plan["id"], priority=10)
|
||||
await db.set_model_route("gpt-3.5-turbo", plan["id"], priority=5)
|
||||
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.get("/api/plans/routes/models")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_model_route(self, temp_db, temp_storage):
|
||||
"""测试删除模型路由"""
|
||||
from app import database as db
|
||||
plan = await db.create_plan(name="Test", provider_name="openai", api_key="sk-test")
|
||||
route = await db.set_model_route("gpt-4", plan["id"], priority=10)
|
||||
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.delete(f"/api/plans/routes/models/{route['id']}")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["ok"] is True
|
||||
|
||||
|
||||
class TestTasksAPI:
|
||||
"""Tasks API 测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_task(self, temp_db, temp_storage):
|
||||
"""测试创建任务"""
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post(
|
||||
"/api/queue",
|
||||
json={
|
||||
"task_type": "image",
|
||||
"request_payload": {"prompt": "a cat"},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "id" in data
|
||||
assert data["status"] == "pending"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_tasks(self, temp_db, temp_storage):
|
||||
"""测试列出任务"""
|
||||
from app import database as db
|
||||
await db.create_task(task_type="image", request_payload={"prompt": "cat"})
|
||||
t = await db.create_task(task_type="voice", request_payload={"text": "hello"})
|
||||
await db.update_task(t["id"], status="running")
|
||||
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.get("/api/queue")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data) == 2
|
||||
|
||||
# 测试过滤
|
||||
response = await client.get("/api/queue?status=pending")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_task(self, temp_db, temp_storage):
|
||||
"""测试获取任务"""
|
||||
from app import database as db
|
||||
task = await db.create_task(task_type="image", request_payload={"prompt": "cat"})
|
||||
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.get(f"/api/queue/{task['id']}")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == task["id"]
|
||||
assert data["task_type"] == "image"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_task(self, temp_db, temp_storage):
|
||||
"""测试取消任务"""
|
||||
from app import database as db
|
||||
task = await db.create_task(task_type="image", request_payload={"prompt": "cat"})
|
||||
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post(f"/api/queue/{task['id']}/cancel")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["ok"] is True
|
||||
|
||||
# 验证状态
|
||||
updated = await db.get_task(task["id"])
|
||||
assert updated["status"] == "cancelled"
|
||||
|
||||
|
||||
class TestProxyAPI:
|
||||
"""代理 API 测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_completions_missing_auth(self, temp_db, temp_storage):
|
||||
"""测试聊天请求缺少鉴权"""
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post(
|
||||
"/v1/chat/completions",
|
||||
json={"model": "gpt-4", "messages": [{"role": "user", "content": "hello"}]},
|
||||
)
|
||||
|
||||
# 默认 proxy_api_key 为空时应该通过，但如果设置了则返回 401
|
||||
# 这里我们测试没有设置 key 的情况
|
||||
# 实际行为取决于配置
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_completions_invalid_model(self, temp_db, temp_storage):
|
||||
"""测试聊天请求无效模型"""
|
||||
from app import database as db
|
||||
|
||||
# 设置空的 proxy_api_key 以跳过鉴权测试
|
||||
import app.config as config_module
|
||||
original_key = config_module.settings.server.proxy_api_key
|
||||
config_module.settings.server.proxy_api_key = ""
|
||||
|
||||
try:
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post(
|
||||
"/v1/chat/completions",
|
||||
json={"model": "nonexistent-model", "messages": [{"role": "user", "content": "hello"}]},
|
||||
)
|
||||
finally:
|
||||
config_module.settings.server.proxy_api_key = original_key
|
||||
|
||||
# 应该返回 404 找不到模型路由
|
||||
assert response.status_code == 404
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_format_conversion(self, temp_db, temp_storage):
|
||||
"""测试 Anthropic 格式转换"""
|
||||
from app import database as db
|
||||
|
||||
import app.config as config_module
|
||||
original_key = config_module.settings.server.proxy_api_key
|
||||
config_module.settings.server.proxy_api_key = ""
|
||||
|
||||
# 创建路由
|
||||
plan = await db.create_plan(name="Test", provider_name="openai", api_key="sk-test")
|
||||
await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule", quota_total=100)
|
||||
await db.set_model_route("gpt-4", plan["id"])
|
||||
|
||||
try:
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post(
|
||||
"/v1/messages",
|
||||
json={
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 100,
|
||||
},
|
||||
)
|
||||
finally:
|
||||
config_module.settings.server.proxy_api_key = original_key
|
||||
|
||||
# 由于没有真实的 API，可能会失败，但至少验证格式转换不抛出错误
|
||||
# 实际会返回 502 或类似的错误，因为没有真实后端
|
||||
@@ -0,0 +1,298 @@
|
||||
"""数据库模块测试"""
|
||||
|
||||
import pytest
|
||||
|
||||
from app import database as db
|
||||
|
||||
|
||||
class TestPlanCRUD:
|
||||
"""Plan CRUD 操作测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_plan(self, db):
|
||||
"""测试创建 Plan"""
|
||||
plan = await db.create_plan(
|
||||
name="Test Plan",
|
||||
provider_name="openai",
|
||||
api_key="sk-test",
|
||||
api_base="https://api.openai.com/v1",
|
||||
plan_type="coding",
|
||||
supported_models=["gpt-4"],
|
||||
)
|
||||
assert "id" in plan
|
||||
assert plan["name"] == "Test Plan"
|
||||
assert plan["provider_name"] == "openai"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_plan(self, db):
|
||||
"""测试获取 Plan"""
|
||||
created = await db.create_plan(
|
||||
name="Get Test",
|
||||
provider_name="openai",
|
||||
api_key="sk-test",
|
||||
)
|
||||
plan = await db.get_plan(created["id"])
|
||||
assert plan is not None
|
||||
assert plan["name"] == "Get Test"
|
||||
assert plan["enabled"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_plan_not_found(self, db):
|
||||
"""测试获取不存在的 Plan"""
|
||||
plan = await db.get_plan("nonexistent")
|
||||
assert plan is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_plans(self, db):
|
||||
"""测试列出 Plans"""
|
||||
await db.create_plan(name="Plan 1", provider_name="openai", api_key="sk1")
|
||||
await db.create_plan(name="Plan 2", provider_name="kimi", api_key="sk2")
|
||||
await db.create_plan(name="Plan 3", provider_name="google", api_key="sk3", enabled=False)
|
||||
|
||||
plans = await db.list_plans()
|
||||
assert len(plans) == 3
|
||||
|
||||
enabled_only = await db.list_plans(enabled_only=True)
|
||||
assert len(enabled_only) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_plan(self, db):
|
||||
"""测试更新 Plan"""
|
||||
plan = await db.create_plan(name="Old Name", provider_name="openai", api_key="sk-test")
|
||||
ok = await db.update_plan(plan["id"], name="New Name", plan_type="token")
|
||||
assert ok is True
|
||||
|
||||
updated = await db.get_plan(plan["id"])
|
||||
assert updated["name"] == "New Name"
|
||||
assert updated["plan_type"] == "token"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_plan(self, db):
|
||||
"""测试删除 Plan"""
|
||||
plan = await db.create_plan(name="To Delete", provider_name="openai", api_key="sk-test")
|
||||
ok = await db.delete_plan(plan["id"])
|
||||
assert ok is True
|
||||
|
||||
deleted = await db.get_plan(plan["id"])
|
||||
assert deleted is None
|
||||
|
||||
|
||||
class TestQuotaRuleCRUD:
|
||||
"""QuotaRule CRUD 操作测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_quota_rule(self, db):
|
||||
"""测试创建 QuotaRule"""
|
||||
plan = await db.create_plan(name="Test Plan", provider_name="openai", api_key="sk-test")
|
||||
rule = await db.create_quota_rule(
|
||||
plan_id=plan["id"],
|
||||
rule_name="Daily Limit",
|
||||
quota_total=100,
|
||||
quota_unit="requests",
|
||||
refresh_type="calendar_cycle",
|
||||
calendar_unit="daily",
|
||||
calendar_anchor={"hour": 0},
|
||||
)
|
||||
assert "id" in rule
|
||||
assert rule["rule_name"] == "Daily Limit"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_quota_rules(self, db):
|
||||
"""测试列出 QuotaRules"""
|
||||
plan = await db.create_plan(name="Test Plan", provider_name="openai", api_key="sk-test")
|
||||
await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule 1", quota_total=100)
|
||||
await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule 2", quota_total=200)
|
||||
|
||||
rules = await db.list_quota_rules(plan["id"])
|
||||
assert len(rules) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_quota_rule(self, db):
|
||||
"""测试更新 QuotaRule"""
|
||||
plan = await db.create_plan(name="Test Plan", provider_name="openai", api_key="sk-test")
|
||||
rule = await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule", quota_total=100)
|
||||
|
||||
ok = await db.update_quota_rule(rule["id"], quota_total=200)
|
||||
assert ok is True
|
||||
|
||||
rules = await db.list_quota_rules(plan["id"])
|
||||
assert rules[0]["quota_total"] == 200
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_quota_rule(self, db):
|
||||
"""测试删除 QuotaRule"""
|
||||
plan = await db.create_plan(name="Test Plan", provider_name="openai", api_key="sk-test")
|
||||
rule = await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule", quota_total=100)
|
||||
|
||||
ok = await db.delete_quota_rule(rule["id"])
|
||||
assert ok is True
|
||||
|
||||
rules = await db.list_quota_rules(plan["id"])
|
||||
assert len(rules) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_increment_quota_used(self, db):
|
||||
"""测试增加已用额度"""
|
||||
plan = await db.create_plan(name="Test Plan", provider_name="openai", api_key="sk-test")
|
||||
await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule", quota_total=100, quota_unit="requests")
|
||||
await db.create_quota_rule(plan_id=plan["id"], rule_name="Token Rule", quota_total=1000, quota_unit="tokens")
|
||||
|
||||
await db.increment_quota_used(plan["id"], token_count=50)
|
||||
|
||||
rules = await db.list_quota_rules(plan["id"])
|
||||
# requests 类型 +1
|
||||
assert rules[0]["quota_used"] == 1
|
||||
# tokens 类型 +50
|
||||
assert rules[1]["quota_used"] == 50
|
||||
|
||||
|
||||
class TestCheckPlanAvailable:
|
||||
"""Plan 可用性检查测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_rules_available(self, db):
|
||||
"""测试所有规则有余量"""
|
||||
plan = await db.create_plan(name="Test Plan", provider_name="openai", api_key="sk-test")
|
||||
rule = await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule", quota_total=100)
|
||||
|
||||
available = await db.check_plan_available(plan["id"])
|
||||
assert available is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_rule_exhausted(self, db):
|
||||
"""测试一个规则耗尽"""
|
||||
plan = await db.create_plan(name="Test Plan", provider_name="openai", api_key="sk-test")
|
||||
rule = await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule", quota_total=100)
|
||||
await db.update_quota_rule(rule["id"], quota_used=100)
|
||||
|
||||
available = await db.check_plan_available(plan["id"])
|
||||
assert available is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_enabled_rules(self, db):
|
||||
"""测试没有启用的规则"""
|
||||
plan = await db.create_plan(name="Test Plan", provider_name="openai", api_key="sk-test")
|
||||
rule = await db.create_quota_rule(
|
||||
plan_id=plan["id"], rule_name="Rule", quota_total=100, enabled=False
|
||||
)
|
||||
|
||||
available = await db.check_plan_available(plan["id"])
|
||||
assert available is False # 当前实现会返回 True，这是需要修复的
|
||||
|
||||
|
||||
class TestModelRoute:
|
||||
"""模型路由测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_model_route(self, db):
|
||||
"""测试设置模型路由"""
|
||||
plan = await db.create_plan(name="Test Plan", provider_name="openai", api_key="sk-test")
|
||||
route = await db.set_model_route("gpt-4", plan["id"], priority=10)
|
||||
|
||||
assert "id" in route
|
||||
assert route["model_name"] == "gpt-4"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_model_route_update_existing(self, db):
|
||||
"""测试设置已存在的路由"""
|
||||
plan = await db.create_plan(name="Test Plan", provider_name="openai", api_key="sk-test")
|
||||
route1 = await db.set_model_route("gpt-4", plan["id"], priority=10)
|
||||
route2 = await db.set_model_route("gpt-4", plan["id"], priority=20)
|
||||
|
||||
# 原实现会创建新 ID，当前实现会更新但返回新 ID
|
||||
# 这里只验证不会抛出错误
|
||||
assert route2["model_name"] == "gpt-4"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_model(self, db):
|
||||
"""测试解析模型路由"""
|
||||
plan1 = await db.create_plan(name="Plan 1", provider_name="openai", api_key="sk1")
|
||||
plan2 = await db.create_plan(name="Plan 2", provider_name="kimi", api_key="sk2")
|
||||
|
||||
await db.set_model_route("gpt-4", plan1["id"], priority=10)
|
||||
await db.set_model_route("gpt-4", plan2["id"], priority=5) # 较低优先级
|
||||
|
||||
# 应该返回高优先级的 plan1
|
||||
resolved = await db.resolve_model("gpt-4")
|
||||
assert resolved == plan1["id"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_model_with_fallback(self, db):
|
||||
"""测试额度耗尽时的 fallback"""
|
||||
plan1 = await db.create_plan(name="Plan 1", provider_name="openai", api_key="sk1")
|
||||
plan2 = await db.create_plan(name="Plan 2", provider_name="kimi", api_key="sk2")
|
||||
|
||||
await db.set_model_route("gpt-4", plan1["id"], priority=10)
|
||||
await db.set_model_route("gpt-4", plan2["id"], priority=5)
|
||||
|
||||
# plan1 额度耗尽
|
||||
rule = await db.create_quota_rule(plan_id=plan1["id"], rule_name="Rule", quota_total=100)
|
||||
await db.update_quota_rule(rule["id"], quota_used=100)
|
||||
|
||||
# 应该 fallback 到 plan2
|
||||
resolved = await db.resolve_model("gpt-4")
|
||||
assert resolved == plan2["id"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_model_route(self, db):
|
||||
"""测试删除模型路由"""
|
||||
plan = await db.create_plan(name="Test Plan", provider_name="openai", api_key="sk-test")
|
||||
route = await db.set_model_route("gpt-4", plan["id"], priority=10)
|
||||
|
||||
ok = await db.delete_model_route(route["id"])
|
||||
assert ok is True
|
||||
|
||||
routes = await db.list_model_routes()
|
||||
assert len(routes) == 0
|
||||
|
||||
|
||||
class TestTaskQueue:
|
||||
"""任务队列测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_task(self, db):
|
||||
"""测试创建任务"""
|
||||
task = await db.create_task(
|
||||
task_type="image",
|
||||
request_payload={"prompt": "A cat"},
|
||||
priority=1,
|
||||
)
|
||||
assert "id" in task
|
||||
assert task["status"] == "pending"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_tasks(self, db):
|
||||
"""测试列出任务"""
|
||||
await db.create_task(task_type="image", request_payload={})
|
||||
t2 = await db.create_task(task_type="voice", request_payload={})
|
||||
await db.update_task(t2["id"], status="running")
|
||||
|
||||
pending = await db.list_tasks(status="pending")
|
||||
assert len(pending) == 1
|
||||
|
||||
all_tasks = await db.list_tasks()
|
||||
assert len(all_tasks) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_task(self, db):
|
||||
"""测试更新任务"""
|
||||
task = await db.create_task(task_type="image", request_payload={})
|
||||
|
||||
ok = await db.update_task(task["id"], status="completed")
|
||||
assert ok is True
|
||||
|
||||
updated = await db.get_task(task["id"])
|
||||
assert updated["status"] == "completed"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_task(self, db):
|
||||
"""测试获取任务"""
|
||||
task = await db.create_task(
|
||||
task_type="image",
|
||||
request_payload={"prompt": "test"},
|
||||
)
|
||||
|
||||
fetched = await db.get_task(task["id"])
|
||||
assert fetched is not None
|
||||
assert fetched["task_type"] == "image"
|
||||
assert fetched["request_payload"] == {"prompt": "test"}
|
||||
@@ -0,0 +1,243 @@
|
||||
"""Provider 模块测试"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from app.providers import ProviderRegistry
|
||||
from app.providers.base import BaseProvider, Capability, QuotaInfo
|
||||
|
||||
|
||||
class MockProvider(BaseProvider):
|
||||
"""测试用 Mock Provider"""
|
||||
name = "mock"
|
||||
display_name = "Mock Provider"
|
||||
capabilities = [Capability.CHAT, Capability.IMAGE]
|
||||
|
||||
async def chat(self, messages, model, plan, stream=True, **kwargs):
|
||||
if stream:
|
||||
yield 'data: {"choices": [{"delta": {"content": "hello"}}]}'
|
||||
else:
|
||||
yield '{"choices": [{"message": {"content": "hello"}}]}'
|
||||
|
||||
async def generate_image(self, prompt, plan, **kwargs):
|
||||
return {"url": f"https://example.com/{prompt}.png"}
|
||||
|
||||
async def query_quota(self, plan):
|
||||
return QuotaInfo(quota_used=50, quota_total=100, quota_remaining=50, unit="tokens")
|
||||
|
||||
|
||||
class TestProviderRegistry:
|
||||
"""Provider 注册表测试"""
|
||||
|
||||
def test_register_and_get(self):
|
||||
"""测试注册和获取 Provider"""
|
||||
mock = MockProvider()
|
||||
ProviderRegistry._providers["mock"] = mock
|
||||
|
||||
retrieved = ProviderRegistry.get("mock")
|
||||
assert retrieved is mock
|
||||
assert retrieved.name == "mock"
|
||||
|
||||
def test_get_nonexistent(self):
|
||||
"""测试获取不存在的 Provider"""
|
||||
retrieved = ProviderRegistry.get("nonexistent")
|
||||
assert retrieved is None
|
||||
|
||||
def test_all_providers(self):
|
||||
"""测试获取所有 Provider"""
|
||||
mock = MockProvider()
|
||||
ProviderRegistry._providers = {"mock": mock}
|
||||
|
||||
all_providers = ProviderRegistry.all()
|
||||
assert "mock" in all_providers
|
||||
assert all_providers["mock"] is mock
|
||||
|
||||
def test_by_capability(self):
|
||||
"""测试按能力筛选 Provider"""
|
||||
mock = MockProvider()
|
||||
ProviderRegistry._providers = {"mock": mock}
|
||||
|
||||
chat_providers = ProviderRegistry.by_capability(Capability.CHAT)
|
||||
assert mock in chat_providers
|
||||
|
||||
video_providers = ProviderRegistry.by_capability(Capability.VIDEO)
|
||||
assert mock not in video_providers
|
||||
|
||||
|
||||
class TestBaseProvider:
|
||||
"""BaseProvider 测试"""
|
||||
|
||||
def test_build_headers(self):
|
||||
"""测试构建请求头"""
|
||||
provider = MockProvider()
|
||||
plan = {
|
||||
"api_key": "sk-test-key",
|
||||
"extra_headers": {"X-Custom": "value"},
|
||||
}
|
||||
|
||||
headers = provider._build_headers(plan)
|
||||
assert headers["Authorization"] == "Bearer sk-test-key"
|
||||
assert headers["X-Custom"] == "value"
|
||||
|
||||
def test_build_headers_no_key(self):
|
||||
"""测试无 API Key 时的请求头"""
|
||||
provider = MockProvider()
|
||||
plan = {}
|
||||
|
||||
headers = provider._build_headers(plan)
|
||||
assert "Authorization" not in headers
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
|
||||
def test_base_url(self):
|
||||
"""测试获取基础 URL"""
|
||||
provider = MockProvider()
|
||||
plan = {"api_base": "https://api.example.com/v1/"}
|
||||
|
||||
url = provider._base_url(plan)
|
||||
assert url == "https://api.example.com/v1"
|
||||
|
||||
def test_base_url_empty(self):
|
||||
"""测试空基础 URL"""
|
||||
provider = MockProvider()
|
||||
plan = {}
|
||||
|
||||
url = provider._base_url(plan)
|
||||
assert url == ""
|
||||
|
||||
|
||||
class TestMockProvider:
|
||||
"""Mock Provider 功能测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_stream(self):
|
||||
"""测试流式聊天"""
|
||||
provider = MockProvider()
|
||||
plan = {"api_key": "sk-test"}
|
||||
|
||||
chunks = []
|
||||
async for chunk in provider.chat([], "gpt-4", plan, stream=True):
|
||||
chunks.append(chunk)
|
||||
|
||||
assert len(chunks) == 1
|
||||
assert "hello" in chunks[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_non_stream(self):
|
||||
"""测试非流式聊天"""
|
||||
provider = MockProvider()
|
||||
plan = {"api_key": "sk-test"}
|
||||
|
||||
chunks = []
|
||||
async for chunk in provider.chat([], "gpt-4", plan, stream=False):
|
||||
chunks.append(chunk)
|
||||
|
||||
assert len(chunks) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_image(self):
|
||||
"""测试图片生成"""
|
||||
provider = MockProvider()
|
||||
plan = {"api_key": "sk-test"}
|
||||
|
||||
result = await provider.generate_image("a cat", plan)
|
||||
assert "url" in result
|
||||
assert "cat" in result["url"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_quota(self):
|
||||
"""测试额度查询"""
|
||||
provider = MockProvider()
|
||||
plan = {"api_key": "sk-test"}
|
||||
|
||||
info = await provider.query_quota(plan)
|
||||
assert info.quota_used == 50
|
||||
assert info.quota_total == 100
|
||||
assert info.quota_remaining == 50
|
||||
assert info.unit == "tokens"
|
||||
|
||||
|
||||
class TestOpenAIProvider:
|
||||
"""OpenAI Provider 测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_requires_api_key(self):
|
||||
"""测试需要 API Key"""
|
||||
from app.providers.openai_provider import OpenAIProvider
|
||||
|
||||
provider = OpenAIProvider()
|
||||
plan = {"api_base": "https://api.openai.com/v1"}
|
||||
|
||||
# 没有 API key 不应该抛出错误，但 headers 中不应该有 Authorization
|
||||
headers = provider._build_headers(plan)
|
||||
assert "Authorization" not in headers or headers["Authorization"] == "Bearer "
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_headers_with_extra(self):
|
||||
"""测试额外的请求头"""
|
||||
from app.providers.openai_provider import OpenAIProvider
|
||||
|
||||
provider = OpenAIProvider()
|
||||
plan = {
|
||||
"api_key": "sk-test",
|
||||
"extra_headers": {"OpenAI-Organization": "org-123"},
|
||||
}
|
||||
|
||||
headers = provider._build_headers(plan)
|
||||
assert headers["Authorization"] == "Bearer sk-test"
|
||||
assert headers["OpenAI-Organization"] == "org-123"
|
||||
|
||||
|
||||
class TestKimiProvider:
|
||||
"""Kimi Provider 测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_metadata(self):
|
||||
"""测试 Provider 元数据"""
|
||||
from app.providers.kimi import KimiProvider
|
||||
|
||||
provider = KimiProvider()
|
||||
assert provider.name == "kimi"
|
||||
assert provider.display_name == "Kimi (Moonshot)"
|
||||
assert Capability.CHAT in provider.capabilities
|
||||
|
||||
|
||||
class TestMiniMaxProvider:
|
||||
"""MiniMax Provider 测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_metadata(self):
|
||||
"""测试 Provider 元数据"""
|
||||
from app.providers.minimax import MiniMaxProvider
|
||||
|
||||
provider = MiniMaxProvider()
|
||||
assert provider.name == "minimax"
|
||||
assert provider.display_name == "MiniMax"
|
||||
assert Capability.CHAT in provider.capabilities
|
||||
|
||||
|
||||
class TestGoogleProvider:
|
||||
"""Google Provider 测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_metadata(self):
|
||||
"""测试 Provider 元数据"""
|
||||
from app.providers.google import GoogleProvider
|
||||
|
||||
provider = GoogleProvider()
|
||||
assert provider.name == "google"
|
||||
assert provider.display_name == "Google Gemini"
|
||||
assert Capability.CHAT in provider.capabilities
|
||||
|
||||
|
||||
class TestZhipuProvider:
|
||||
"""智谱 Provider 测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_metadata(self):
|
||||
"""测试 Provider 元数据"""
|
||||
from app.providers.zhipu import ZhipuProvider
|
||||
|
||||
provider = ZhipuProvider()
|
||||
assert provider.name == "zhipu"
|
||||
assert "GLM" in provider.display_name or "智谱" in provider.display_name
|
||||
assert Capability.CHAT in provider.capabilities
|
||||
@@ -0,0 +1,190 @@
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
|
||||
from app import database as db
|
||||
from app.config import settings
|
||||
from app.routers import proxy
|
||||
from app.services import scheduler
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolated_db(tmp_path):
|
||||
old_path = settings.database.path
|
||||
settings.database.path = str(tmp_path / "test.db")
|
||||
asyncio.run(db.close_db())
|
||||
yield
|
||||
asyncio.run(db.close_db())
|
||||
settings.database.path = old_path
|
||||
|
||||
|
||||
def test_update_functions_report_missing_rows():
|
||||
# update_plan 现在返回 cur.rowcount > 0，不存在的行返回 False
|
||||
result = asyncio.run(db.update_plan("missing", name="x"))
|
||||
assert result is False
|
||||
assert asyncio.run(db.update_quota_rule("missing", rule_name="x")) is False
|
||||
assert asyncio.run(db.update_task("missing", status="cancelled")) is False
|
||||
|
||||
|
||||
def test_update_functions_report_success_for_existing_rows():
|
||||
plan = asyncio.run(db.create_plan(name="a", provider_name="openai"))
|
||||
rule = asyncio.run(
|
||||
db.create_quota_rule(plan_id=plan["id"], rule_name="r", quota_total=10)
|
||||
)
|
||||
task = asyncio.run(db.create_task(task_type="image", request_payload={"prompt": "p"}))
|
||||
assert asyncio.run(db.update_plan(plan["id"], name="b")) is True
|
||||
assert asyncio.run(db.update_quota_rule(rule["id"], quota_total=20)) is True
|
||||
assert asyncio.run(db.update_task(task["id"], status="running")) is True
|
||||
|
||||
|
||||
def test_execute_task_avoids_duplicate_prompt_argument():
|
||||
class Provider:
|
||||
async def generate_image(self, prompt: str, plan: dict, **kwargs):
|
||||
return {"prompt": prompt, "kwargs": kwargs}
|
||||
|
||||
task = {
|
||||
"id": "t1",
|
||||
"task_type": "image",
|
||||
"request_payload": {"prompt": "hello", "model": "m1"},
|
||||
}
|
||||
result = asyncio.run(scheduler._execute_task(Provider(), {}, task))
|
||||
assert result["prompt"] == "hello"
|
||||
assert result["kwargs"] == {"model": "m1"}
|
||||
|
||||
|
||||
def test_compute_next_calendar_monthly_handles_large_day_anchor():
|
||||
after = datetime(2026, 2, 15, 12, 0, tzinfo=timezone.utc)
|
||||
next_at = scheduler._compute_next_calendar("monthly", {"day": 31, "hour": 0}, after)
|
||||
assert next_at == datetime(2026, 2, 28, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
after2 = datetime(2026, 2, 28, 1, 0, tzinfo=timezone.utc)
|
||||
next_at2 = scheduler._compute_next_calendar("monthly", {"day": 31, "hour": 0}, after2)
|
||||
assert next_at2 == datetime(2026, 3, 31, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def test_refresh_quota_rules_isolates_api_sync_errors(monkeypatch):
|
||||
fixed_now = datetime(2026, 3, 1, 12, 0, tzinfo=timezone.utc)
|
||||
updates = []
|
||||
rules = [
|
||||
{
|
||||
"id": "r1",
|
||||
"rule_name": "api",
|
||||
"refresh_type": "api_sync",
|
||||
"plan_id": "p1",
|
||||
"last_refresh_at": None,
|
||||
"next_refresh_at": None,
|
||||
},
|
||||
{
|
||||
"id": "r2",
|
||||
"rule_name": "fixed",
|
||||
"refresh_type": "fixed_interval",
|
||||
"interval_hours": 1,
|
||||
"last_refresh_at": (fixed_now - timedelta(hours=2)).isoformat(),
|
||||
"next_refresh_at": (fixed_now - timedelta(minutes=1)).isoformat(),
|
||||
},
|
||||
]
|
||||
|
||||
async def fake_get_all_quota_rules():
|
||||
return rules
|
||||
|
||||
async def fake_get_plan(plan_id: str):
|
||||
return {"id": plan_id, "provider_name": "broken"}
|
||||
|
||||
async def fake_update_quota_rule(rule_id: str, **fields):
|
||||
updates.append((rule_id, fields))
|
||||
return True
|
||||
|
||||
class BadProvider:
|
||||
async def query_quota(self, plan: dict):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
import app.providers as providers
|
||||
|
||||
monkeypatch.setattr(scheduler, "_now", lambda: fixed_now)
|
||||
monkeypatch.setattr(scheduler.db, "get_all_quota_rules", fake_get_all_quota_rules)
|
||||
monkeypatch.setattr(scheduler.db, "get_plan", fake_get_plan)
|
||||
monkeypatch.setattr(scheduler.db, "update_quota_rule", fake_update_quota_rule)
|
||||
monkeypatch.setattr(
|
||||
providers.ProviderRegistry,
|
||||
"get",
|
||||
classmethod(lambda cls, name: BadProvider()),
|
||||
)
|
||||
|
||||
asyncio.run(scheduler._refresh_quota_rules())
|
||||
assert any(rule_id == "r2" and fields.get("quota_used") == 0 for rule_id, fields in updates)
|
||||
|
||||
|
||||
def test_anthropic_route_forwards_extra_kwargs(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
class Provider:
|
||||
async def chat(self, messages, model, plan, stream=True, **kwargs):
|
||||
captured["kwargs"] = kwargs
|
||||
payload = {
|
||||
"choices": [{"message": {"content": "ok"}}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3},
|
||||
}
|
||||
yield json.dumps(payload)
|
||||
|
||||
async def fake_resolve_model(model: str):
|
||||
return "p1"
|
||||
|
||||
async def fake_get_plan(plan_id: str):
|
||||
return {"id": plan_id, "name": "p", "provider_name": "x"}
|
||||
|
||||
async def fake_check_plan_available(plan_id: str):
|
||||
return True
|
||||
|
||||
async def fake_increment_quota_used(plan_id: str, token_count: int = 0):
|
||||
return None
|
||||
|
||||
provider = Provider()
|
||||
monkeypatch.setattr(proxy.db, "resolve_model", fake_resolve_model)
|
||||
monkeypatch.setattr(proxy.db, "get_plan", fake_get_plan)
|
||||
monkeypatch.setattr(proxy.db, "check_plan_available", fake_check_plan_available)
|
||||
monkeypatch.setattr(proxy.db, "increment_quota_used", fake_increment_quota_used)
|
||||
monkeypatch.setattr(
|
||||
proxy.ProviderRegistry,
|
||||
"get",
|
||||
classmethod(lambda cls, name: provider),
|
||||
)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(proxy.router)
|
||||
old_key = settings.server.proxy_api_key
|
||||
settings.server.proxy_api_key = ""
|
||||
try:
|
||||
client = TestClient(app)
|
||||
resp = client.post(
|
||||
"/v1/messages",
|
||||
json={
|
||||
"model": "m1",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"temperature": 0.2,
|
||||
"max_tokens": 88,
|
||||
},
|
||||
)
|
||||
finally:
|
||||
settings.server.proxy_api_key = old_key
|
||||
|
||||
assert resp.status_code == 200
|
||||
assert captured["kwargs"] == {"temperature": 0.2, "max_tokens": 88}
|
||||
|
||||
|
||||
def test_verify_key_requires_auth_when_using_default_value():
|
||||
old_key = settings.server.proxy_api_key
|
||||
settings.server.proxy_api_key = "sk-plan-manage-change-me"
|
||||
try:
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
proxy._verify_key(None)
|
||||
finally:
|
||||
settings.server.proxy_api_key = old_key
|
||||
assert exc.value.status_code == 401
|
||||
@@ -0,0 +1,347 @@
|
||||
"""调度器模块测试"""
|
||||
|
||||
import pytest
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from app.services.scheduler import (
|
||||
_now,
|
||||
_parse_dt,
|
||||
_compute_next_calendar,
|
||||
_refresh_quota_rules,
|
||||
_process_task_queue,
|
||||
_execute_task,
|
||||
)
|
||||
|
||||
|
||||
class TestUtilityFunctions:
|
||||
"""工具函数测试"""
|
||||
|
||||
def test_now_returns_utc(self):
|
||||
"""测试 _now 返回 UTC 时间"""
|
||||
now = _now()
|
||||
assert now.tzinfo == timezone.utc
|
||||
|
||||
def test_parse_dt_valid(self):
|
||||
"""测试解析有效的日期时间字符串"""
|
||||
s = "2026-03-31T12:00:00+00:00"
|
||||
result = _parse_dt(s)
|
||||
assert result is not None
|
||||
assert result.year == 2026
|
||||
assert result.month == 3
|
||||
assert result.day == 31
|
||||
|
||||
def test_parse_dt_none(self):
|
||||
"""测试解析 None"""
|
||||
result = _parse_dt(None)
|
||||
assert result is None
|
||||
|
||||
def test_parse_dt_invalid(self):
|
||||
"""测试解析无效字符串"""
|
||||
result = _parse_dt("not-a-date")
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestComputeNextCalendar:
|
||||
"""计算下一个自然周期测试"""
|
||||
|
||||
def test_daily_anchor(self):
|
||||
"""测试每日刷新点"""
|
||||
# 假设现在是 3月31日 15:00
|
||||
after = datetime(2026, 3, 31, 15, 0, tzinfo=timezone.utc)
|
||||
anchor = {"hour": 0}
|
||||
|
||||
next_time = _compute_next_calendar("daily", anchor, after)
|
||||
|
||||
# 应该是第二天 0 点
|
||||
assert next_time.day == 1
|
||||
assert next_time.hour == 0
|
||||
assert next_time.month == 4
|
||||
|
||||
def test_daily_anchor_before_hour(self):
|
||||
"""测试每日刷新点（当前时间在刷新点之前）"""
|
||||
after = datetime(2026, 3, 31, 8, 0, tzinfo=timezone.utc)
|
||||
anchor = {"hour": 10}
|
||||
|
||||
next_time = _compute_next_calendar("daily", anchor, after)
|
||||
|
||||
# 应该是当天 10 点
|
||||
assert next_time.day == 31
|
||||
assert next_time.hour == 10
|
||||
|
||||
def test_weekly_anchor(self):
|
||||
"""测试每周刷新点（周一 0 点）"""
|
||||
# 假设现在是周三 (2026-03-31 是周二，用 4月1日周三)
|
||||
after = datetime(2026, 4, 1, 10, 0, tzinfo=timezone.utc) # 周三
|
||||
anchor = {"weekday": 1, "hour": 0} # 周一
|
||||
|
||||
next_time = _compute_next_calendar("weekly", anchor, after)
|
||||
|
||||
# 应该是下周周一
|
||||
assert next_time.hour == 0
|
||||
# 4月1日是周三，下一个周一是4月6日
|
||||
assert next_time.day == 6
|
||||
|
||||
def test_weekly_anchor_same_day(self):
|
||||
"""测试每周刷新点（当天是刷新日但已过时间）"""
|
||||
after = datetime(2026, 3, 31, 10, 0, tzinfo=timezone.utc) # 周二
|
||||
anchor = {"weekday": 2, "hour": 0} # 周二
|
||||
|
||||
next_time = _compute_next_calendar("weekly", anchor, after)
|
||||
|
||||
# 应该是下周二
|
||||
assert next_time.day == 7 # 3月31日是周二，下周二是4月7日
|
||||
|
||||
def test_monthly_anchor_first_day(self):
|
||||
"""测试每月刷新点（每月1号）"""
|
||||
after = datetime(2026, 3, 15, 10, 0, tzinfo=timezone.utc)
|
||||
anchor = {"day": 1, "hour": 0}
|
||||
|
||||
next_time = _compute_next_calendar("monthly", anchor, after)
|
||||
|
||||
# 应该是4月1号
|
||||
assert next_time.day == 1
|
||||
assert next_time.month == 4
|
||||
assert next_time.hour == 0
|
||||
|
||||
def test_monthly_anchor_last_day(self):
|
||||
"""测试每月刷新点（31号）"""
|
||||
after = datetime(2026, 3, 15, 10, 0, tzinfo=timezone.utc)
|
||||
anchor = {"day": 31, "hour": 0}
|
||||
|
||||
next_time = _compute_next_calendar("monthly", anchor, after)
|
||||
|
||||
# 3月31号
|
||||
assert next_time.day == 31
|
||||
assert next_time.month == 3
|
||||
|
||||
def test_monthly_anchor_invalid_day(self):
|
||||
"""测试每月刷新点（不存在日期，如2月31号）"""
|
||||
after = datetime(2026, 3, 15, 10, 0, tzinfo=timezone.utc)
|
||||
anchor = {"day": 30, "hour": 0} # 2月没有30号
|
||||
|
||||
next_time = _compute_next_calendar("monthly", anchor, after)
|
||||
|
||||
# 应该是3月30号
|
||||
assert next_time.day == 30
|
||||
assert next_time.month == 3
|
||||
|
||||
|
||||
class TestRefreshQuotaRules:
|
||||
"""额度刷新测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manual_rule_no_refresh(self, temp_db):
|
||||
"""测试手动规则不刷新"""
|
||||
await _refresh_quota_rules()
|
||||
# 手动规则应该被跳过，不抛出错误
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fixed_interval_refresh(self, temp_db):
|
||||
"""测试固定间隔刷新"""
|
||||
from app import database as db
|
||||
from datetime import datetime, timezone, timedelta
|
||||
|
||||
# 创建 Plan 和 Rule
|
||||
plan = await db.create_plan(name="Test", provider_name="openai", api_key="sk-test")
|
||||
rule = await db.create_quota_rule(
|
||||
plan_id=plan["id"],
|
||||
rule_name="Hourly",
|
||||
quota_total=100,
|
||||
refresh_type="fixed_interval",
|
||||
interval_hours=1,
|
||||
)
|
||||
|
||||
# 设置 last_refresh_at 和 next_refresh_at 为过去时间
|
||||
past = datetime.now(timezone.utc) - timedelta(hours=2)
|
||||
await db.update_quota_rule(
|
||||
rule["id"],
|
||||
last_refresh_at=past.isoformat(),
|
||||
next_refresh_at=past.isoformat(),
|
||||
)
|
||||
await db.update_quota_rule(rule["id"], quota_used=99)
|
||||
|
||||
# 执行刷新
|
||||
await _refresh_quota_rules()
|
||||
|
||||
# 验证额度已重置
|
||||
rules = await db.list_quota_rules(plan["id"])
|
||||
assert rules[0]["quota_used"] == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_calendar_cycle_refresh(self, temp_db):
|
||||
"""测试日历周期刷新"""
|
||||
from app import database as db
|
||||
from datetime import datetime, timezone, timedelta
|
||||
|
||||
plan = await db.create_plan(name="Test", provider_name="openai", api_key="sk-test")
|
||||
rule = await db.create_quota_rule(
|
||||
plan_id=plan["id"],
|
||||
rule_name="Daily",
|
||||
quota_total=100,
|
||||
refresh_type="calendar_cycle",
|
||||
calendar_unit="daily",
|
||||
calendar_anchor={"hour": 0},
|
||||
)
|
||||
|
||||
# 设置为昨天
|
||||
past = datetime.now(timezone.utc) - timedelta(days=1)
|
||||
await db.update_quota_rule(
|
||||
rule["id"],
|
||||
last_refresh_at=past.isoformat(),
|
||||
next_refresh_at=past.isoformat(),
|
||||
)
|
||||
await db.update_quota_rule(rule["id"], quota_used=99)
|
||||
|
||||
await _refresh_quota_rules()
|
||||
|
||||
rules = await db.list_quota_rules(plan["id"])
|
||||
assert rules[0]["quota_used"] == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_sync(self, temp_db):
|
||||
"""测试 API 同步刷新"""
|
||||
from app import database as db
|
||||
from app.providers import ProviderRegistry
|
||||
from app.providers.base import QuotaInfo
|
||||
|
||||
# 注册 Mock Provider
|
||||
class MockSyncProvider:
|
||||
name = "mock_sync"
|
||||
capabilities = []
|
||||
|
||||
async def query_quota(self, plan):
|
||||
return QuotaInfo(quota_used=42, quota_total=100, quota_remaining=58, unit="tokens")
|
||||
|
||||
ProviderRegistry._providers["mock_sync"] = MockSyncProvider()
|
||||
|
||||
plan = await db.create_plan(name="Test", provider_name="mock_sync", api_key="sk-test")
|
||||
rule = await db.create_quota_rule(
|
||||
plan_id=plan["id"],
|
||||
rule_name="API Sync",
|
||||
quota_total=100,
|
||||
refresh_type="api_sync",
|
||||
)
|
||||
|
||||
# 设置 last_refresh_at 为超过10分钟前
|
||||
past = datetime.now(timezone.utc) - timedelta(seconds=601)
|
||||
await db.update_quota_rule(rule["id"], last_refresh_at=past.isoformat())
|
||||
|
||||
await _refresh_quota_rules()
|
||||
|
||||
# 验证不会抛出错误即可（实际更新取决于调度逻辑）
|
||||
rules = await db.list_quota_rules(plan["id"])
|
||||
assert rules[0]["rule_name"] == "API Sync"
|
||||
|
||||
|
||||
class TestProcessTaskQueue:
|
||||
"""任务队列处理测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_pending_tasks(self, temp_db):
|
||||
"""测试处理待处理任务"""
|
||||
from app import database as db
|
||||
|
||||
# 创建 Plan
|
||||
plan = await db.create_plan(name="Test", provider_name="openai", api_key="sk-test")
|
||||
await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule", quota_total=100)
|
||||
|
||||
# 创建任务
|
||||
task = await db.create_task(
|
||||
task_type="image",
|
||||
request_payload={"prompt": "test"},
|
||||
plan_id=plan["id"],
|
||||
)
|
||||
|
||||
# Mock Provider
|
||||
from app.providers import ProviderRegistry
|
||||
|
||||
class MockTaskProvider:
|
||||
name = "mock_task"
|
||||
capabilities = []
|
||||
|
||||
async def generate_image(self, prompt, plan, **kwargs):
|
||||
return {"url": "https://example.com/test.png"}
|
||||
|
||||
ProviderRegistry._providers["openai"] = MockTaskProvider()
|
||||
|
||||
await _process_task_queue()
|
||||
|
||||
# 验证任务状态
|
||||
updated_task = await db.get_task(task["id"])
|
||||
# 由于我们没有真实的 Provider 实现完整逻辑，任务可能保持 pending 或失败
|
||||
# 这里只验证不会抛出错误
|
||||
assert updated_task is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skip_task_no_quota(self, temp_db):
|
||||
"""测试跳过额度不足的任务"""
|
||||
from app import database as db
|
||||
|
||||
plan = await db.create_plan(name="Test", provider_name="openai", api_key="sk-test")
|
||||
rule = await db.create_quota_rule(plan_id=plan["id"], rule_name="Rule", quota_total=100)
|
||||
await db.update_quota_rule(rule["id"], quota_used=100)
|
||||
|
||||
await db.create_task(
|
||||
task_type="image",
|
||||
request_payload={"prompt": "test"},
|
||||
plan_id=plan["id"],
|
||||
)
|
||||
|
||||
# 应该跳过而不抛出错误
|
||||
await _process_task_queue()
|
||||
|
||||
|
||||
class TestExecuteTask:
|
||||
"""任务执行测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_image_task(self):
|
||||
"""测试执行图片生成任务"""
|
||||
from app.providers import ProviderRegistry
|
||||
|
||||
class MockImageProvider:
|
||||
name = "mock_img"
|
||||
capabilities = []
|
||||
|
||||
async def generate_image(self, prompt, plan, **kwargs):
|
||||
return {"url": f"https://example.com/{prompt}.png"}
|
||||
|
||||
ProviderRegistry._providers["mock_img"] = MockImageProvider()
|
||||
|
||||
plan = {"api_key": "sk-test"}
|
||||
task = {
|
||||
"id": "task123",
|
||||
"task_type": "image",
|
||||
"request_payload": {"prompt": "a cat", "size": "512x512"},
|
||||
}
|
||||
|
||||
provider = MockImageProvider()
|
||||
result = await _execute_task(provider, plan, task)
|
||||
|
||||
assert "url" in result
|
||||
assert "cat" in result["url"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_unknown_task(self):
|
||||
"""测试执行未知类型任务"""
|
||||
from app.providers import ProviderRegistry
|
||||
|
||||
class MockProvider:
|
||||
name = "mock"
|
||||
capabilities = []
|
||||
|
||||
ProviderRegistry._providers["mock"] = MockProvider()
|
||||
|
||||
plan = {"api_key": "sk-test"}
|
||||
task = {
|
||||
"id": "task123",
|
||||
"task_type": "unknown_type",
|
||||
"request_payload": {},
|
||||
}
|
||||
|
||||
provider = MockProvider()
|
||||
result = await _execute_task(provider, plan, task)
|
||||
|
||||
assert "error" in result
|
||||
assert "Unknown task type" in result["error"]
|
||||
Reference in New Issue
Block a user