54 lines
2.6 KiB
Python
54 lines
2.6 KiB
Python
|
import unittest
|
||
|
|
|
||
|
|
from sqlmodel import SQLModel, Session, create_engine, select
|
||
|
|
|
||
|
|
from app.auth import required_permission
|
||
|
|
from app.catalog import resolve_catalog_model, sync_provider_catalog
|
||
|
|
from app.main import list_available_models, update_model_catalog
|
||
|
|
from app.models import ModelCatalog, ModelProvider
|
||
|
|
|
||
|
|
|
||
|
|
class ModelCatalogTests(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
self.engine = create_engine("sqlite://")
|
||
|
|
SQLModel.metadata.create_all(self.engine)
|
||
|
|
self.db = Session(self.engine)
|
||
|
|
|
||
|
|
def tearDown(self):
|
||
|
|
self.db.close()
|
||
|
|
|
||
|
|
def test_provider_models_sync_to_user_catalog_and_stale_entries_disable(self):
|
||
|
|
provider = ModelProvider(name="目录 Provider", base_url="https://gateway.example/v1",
|
||
|
|
models=["model-a", "model-b"], default_model="model-a")
|
||
|
|
self.db.add(provider); self.db.commit(); self.db.refresh(provider)
|
||
|
|
sync_provider_catalog(self.db, provider); self.db.commit()
|
||
|
|
rows = list_available_models(self.db)
|
||
|
|
self.assertEqual([row["model_name"] for row in rows], ["model-a", "model-b"])
|
||
|
|
first = self.db.exec(
|
||
|
|
select(ModelCatalog).where(ModelCatalog.model_name == "model-a")
|
||
|
|
).one()
|
||
|
|
update_model_catalog(first.id, {"capability_status": "supported", "capabilities": ["chat"]}, self.db)
|
||
|
|
self.assertEqual(self.db.get(ModelCatalog, first.id).capability_status, "supported")
|
||
|
|
provider.models = ["model-a"]; sync_provider_catalog(self.db, provider); self.db.commit()
|
||
|
|
self.assertEqual(len(list_available_models(self.db)), 1)
|
||
|
|
self.assertFalse(self.db.exec(
|
||
|
|
select(ModelCatalog).where(ModelCatalog.model_name == "model-b")
|
||
|
|
).one().enabled)
|
||
|
|
|
||
|
|
def test_catalog_resolution_and_provider_routes_are_separate(self):
|
||
|
|
provider = ModelProvider(name="解析 Provider", base_url="https://gateway.example/v1",
|
||
|
|
models=["model-a"], default_model="model-a")
|
||
|
|
self.db.add(provider); self.db.commit(); self.db.refresh(provider)
|
||
|
|
sync_provider_catalog(self.db, provider); self.db.commit()
|
||
|
|
catalog = self.db.exec(select(ModelCatalog)).one()
|
||
|
|
resolved, resolved_provider = resolve_catalog_model(self.db, catalog.id)
|
||
|
|
self.assertEqual(resolved.model_name, "model-a")
|
||
|
|
self.assertEqual(resolved_provider.id, provider.id)
|
||
|
|
self.assertEqual(required_permission("GET", "/api/providers"), "read")
|
||
|
|
self.assertEqual(required_permission("GET", "/api/models/available"), "read")
|
||
|
|
self.assertEqual(required_permission("GET", "/api/admin/models"), "audit:manage")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|