104 lines
3.9 KiB
Python
104 lines
3.9 KiB
Python
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
|
|
from app.api.model_routing import ModelRoutingUpdateRequest, get_model_routing, update_model_routing
|
|
from app.models import Base, User
|
|
|
|
|
|
class ModelRoutingApiTest(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self) -> None:
|
|
tmp_db = tempfile.NamedTemporaryFile(prefix="pmflow-model-routing-", suffix=".db", delete=False)
|
|
tmp_db.close()
|
|
self.db_path = tmp_db.name
|
|
self.engine = create_async_engine(f"sqlite+aiosqlite:///{self.db_path}", future=True)
|
|
self.session_factory = async_sessionmaker(self.engine, expire_on_commit=False, class_=AsyncSession)
|
|
async with self.engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
tmp_rules = tempfile.NamedTemporaryFile(prefix="pmflow-routing-rules-", suffix=".json", delete=False)
|
|
tmp_rules.write(
|
|
json.dumps(
|
|
{
|
|
"rules": [
|
|
{
|
|
"name": "planner-default",
|
|
"enabled": True,
|
|
"priority": 100,
|
|
"tier": "high",
|
|
"provider": "openai",
|
|
"model": "GLM-5",
|
|
"api_env": "OPENAI_API_KEY",
|
|
"function_exact": "signal_planner_analyze",
|
|
}
|
|
]
|
|
}
|
|
).encode("utf-8")
|
|
)
|
|
tmp_rules.flush()
|
|
tmp_rules.close()
|
|
self.rules_path = tmp_rules.name
|
|
os.environ["OPENCLAW_ROUTING_RULES_PATH"] = self.rules_path
|
|
|
|
async def asyncTearDown(self) -> None:
|
|
await self.engine.dispose()
|
|
for path in (getattr(self, "db_path", None), getattr(self, "rules_path", None)):
|
|
if path:
|
|
try:
|
|
os.remove(path)
|
|
except FileNotFoundError:
|
|
pass
|
|
os.environ.pop("OPENCLAW_ROUTING_RULES_PATH", None)
|
|
|
|
async def test_get_and_update_model_routing(self) -> None:
|
|
async with self.session_factory() as db:
|
|
user = User(
|
|
tenant_id=1,
|
|
username="routing_admin",
|
|
email="routing_admin@example.com",
|
|
password_hash="hash",
|
|
name="路由管理员",
|
|
role="BOSS",
|
|
status="active",
|
|
is_active=True,
|
|
)
|
|
db.add(user)
|
|
await db.commit()
|
|
await db.refresh(user)
|
|
|
|
current = await get_model_routing(
|
|
function_name="signal_planner_analyze",
|
|
db=db,
|
|
current_user=user,
|
|
)
|
|
self.assertEqual(current["code"], 0)
|
|
self.assertEqual(current["data"]["source"], "default")
|
|
|
|
updated = await update_model_routing(
|
|
function_name="signal_planner_analyze",
|
|
body=ModelRoutingUpdateRequest(
|
|
provider="glm",
|
|
model="GLM-4.7",
|
|
api_env="GLM_API_KEY",
|
|
tier="low",
|
|
),
|
|
db=db,
|
|
current_user=user,
|
|
)
|
|
self.assertEqual(updated["code"], 0)
|
|
self.assertEqual(updated["data"]["source"], "override")
|
|
self.assertEqual(updated["data"]["provider"], "glm")
|
|
self.assertEqual(updated["data"]["model"], "GLM-4.7")
|
|
|
|
refreshed = await get_model_routing(
|
|
function_name="signal_planner_analyze",
|
|
db=db,
|
|
current_user=user,
|
|
)
|
|
self.assertEqual(refreshed["code"], 0)
|
|
self.assertEqual(refreshed["data"]["source"], "override")
|
|
self.assertEqual(refreshed["data"]["tier"], "low")
|