pixel/backend/tests/test_models_api_format.py

"""
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