""" Tests for the models API grouped format """ import pytest from fastapi.testclient import TestClient from src.main import app from src.services.provider.registry import ModelRegistry, ModelType, ServiceConfig, ServiceFactory class MockImageService: """Mock image service for testing""" def __init__(self, **kwargs): self.config = kwargs class MockVideoService: """Mock video service for testing""" def __init__(self, **kwargs): self.config = kwargs @pytest.fixture def setup_test_models(): """Setup test models in registry""" # Register test image models image_config_1 = ServiceConfig( id="dashscope/qwen-image", module="test", class_name="MockImageService", name="Qwen Image", type="image", provider="dashscope", model_key="qwen-image", enabled=True, is_default=True, capabilities={"supportsLora": True, "supportsRefImage": True}, resolutions={"1K": {"16:9": "1280*720", "1:1": "1024*1024"}} ) factory_1 = ServiceFactory(image_config_1, MockImageService) ModelRegistry.register_factory("dashscope/qwen-image", factory_1, ModelType.IMAGE, is_default=True) image_config_2 = ServiceConfig( id="modelscope/qwen-image", module="test", class_name="MockImageService", name="ModelScope Qwen Image", type="image", provider="modelscope", model_key="qwen-image", enabled=True, is_default=False ) factory_2 = ServiceFactory(image_config_2, MockImageService) ModelRegistry.register_factory("modelscope/qwen-image", factory_2, ModelType.IMAGE) # Register test video model video_config = ServiceConfig( id="dashscope/wan2.6-video", module="test", class_name="MockVideoService", name="Wan 2.6 Video", type="video", provider="dashscope", model_key="wan2.6-video", enabled=True, is_default=True ) factory_3 = ServiceFactory(video_config, MockVideoService) ModelRegistry.register_factory("dashscope/wan2.6-video", factory_3, ModelType.VIDEO, is_default=True) yield # Cleanup (registry is a singleton, so we need to clean up) # Note: In real tests, you might want to reset the registry def test_models_api_returns_grouped_format(setup_test_models): """Test that /config/models returns grouped HashMap format""" client = TestClient(app) response = client.get("/api/v1/config/models") assert response.status_code == 200 data = response.json() # Check response structure assert "code" in data assert "data" in data assert data["code"] == "0000" # API uses "0000" for success # Check grouped structure grouped_data = data["data"] assert "image" in grouped_data assert "video" in grouped_data assert "audio" in grouped_data assert "llm" in grouped_data # Check that image and video are dicts (HashMaps) assert isinstance(grouped_data["image"], dict) assert isinstance(grouped_data["video"], dict) def test_models_api_image_models_structure(setup_test_models): """Test that image models have correct structure""" client = TestClient(app) response = client.get("/api/v1/config/models") data = response.json() image_models = data["data"]["image"] # Check that we have the expected models assert "dashscope/qwen-image" in image_models assert "modelscope/qwen-image" in image_models # Check dashscope/qwen-image structure qwen_model = image_models["dashscope/qwen-image"] assert qwen_model["id"] == "dashscope/qwen-image" assert qwen_model["name"] == "Qwen Image" assert qwen_model["type"] == "image" assert qwen_model["provider"] == "dashscope" assert qwen_model["model_key"] == "qwen-image" assert qwen_model["is_default"] is True assert qwen_model["enabled"] is True # Check capabilities assert "capabilities" in qwen_model assert qwen_model["capabilities"]["supportsLora"] is True assert qwen_model["capabilities"]["supportsRefImage"] is True # Check resolutions assert "resolutions" in qwen_model assert "1K" in qwen_model["resolutions"] def test_models_api_video_models_structure(setup_test_models): """Test that video models have correct structure""" client = TestClient(app) response = client.get("/api/v1/config/models") data = response.json() video_models = data["data"]["video"] # Check that we have the expected model assert "dashscope/wan2.6-video" in video_models # Check structure wan_model = video_models["dashscope/wan2.6-video"] assert wan_model["id"] == "dashscope/wan2.6-video" assert wan_model["name"] == "Wan 2.6 Video" assert wan_model["type"] == "video" assert wan_model["provider"] == "dashscope" assert wan_model["model_key"] == "wan2.6-video" assert wan_model["is_default"] is True def test_models_api_hashmap_key_matches_id(setup_test_models): """Test that HashMap keys match the id field of each model""" client = TestClient(app) response = client.get("/api/v1/config/models") data = response.json() grouped_data = data["data"] # Check all model types for model_type in ["image", "video", "audio", "llm"]: models = grouped_data[model_type] for model_id, model_config in models.items(): # HashMap key should match the id field assert model_id == model_config["id"], \ f"HashMap key '{model_id}' does not match id field '{model_config['id']}'" def test_models_api_default_flag(setup_test_models): """Test that is_default flag is correctly set""" client = TestClient(app) response = client.get("/api/v1/config/models") data = response.json() image_models = data["data"]["image"] # dashscope/qwen-image should be default assert image_models["dashscope/qwen-image"]["is_default"] is True # modelscope/qwen-image should not be default assert image_models["modelscope/qwen-image"]["is_default"] is False