"""Unit tests for the model catalog resource."""
from Imagen_MCP.resources.models import (
get_model_catalog,
get_model_info,
list_model_ids,
)
class TestGetModelCatalog:
"""Tests for get_model_catalog function."""
def test_get_model_catalog_returns_dict(self):
"""get_model_catalog should return a dictionary."""
catalog = get_model_catalog()
assert isinstance(catalog, dict)
assert "models" in catalog
assert "default_model" in catalog
assert "total_count" in catalog
def test_get_model_catalog_contains_all_models(self):
"""get_model_catalog should contain all available models."""
catalog = get_model_catalog()
assert catalog["total_count"] == 5
assert len(catalog["models"]) == 5
model_ids = [m["id"] for m in catalog["models"]]
assert "imagen-4" in model_ids
assert "imagen-4-fast" in model_ids
assert "imagen-4-ultra" in model_ids
assert "dall-e-3" in model_ids
assert "gpt-image-1" in model_ids
def test_get_model_catalog_default_model(self):
"""get_model_catalog should specify the default model."""
catalog = get_model_catalog()
assert catalog["default_model"] == "imagen-4"
def test_get_model_catalog_model_structure(self):
"""Each model in catalog should have complete structure."""
catalog = get_model_catalog()
for model in catalog["models"]:
assert "id" in model
assert "name" in model
assert "provider" in model
assert "description" in model
assert "capabilities" in model
assert "rate_limit" in model
assert "use_cases" in model
assert "strengths" in model
assert "weaknesses" in model
# Check capabilities structure
caps = model["capabilities"]
assert "max_images_per_request" in caps
assert "supported_sizes" in caps
assert "supports_hd_quality" in caps
assert "supports_style_parameter" in caps
assert "max_prompt_length" in caps
# Check rate limit structure
rate = model["rate_limit"]
assert "messages_per_period" in rate
assert "period_hours" in rate
assert "category" in rate
class TestGetModelInfo:
"""Tests for get_model_info function."""
def test_get_model_info_existing_model(self):
"""get_model_info should return info for existing model."""
info = get_model_info("imagen-4")
assert info is not None
assert info["id"] == "imagen-4"
assert info["name"] == "Imagen 4"
assert info["provider"] == "Google"
def test_get_model_info_nonexistent_model(self):
"""get_model_info should return None for non-existent model."""
info = get_model_info("nonexistent-model")
assert info is None
def test_get_model_info_all_models(self):
"""get_model_info should work for all models."""
model_ids = [
"imagen-4",
"imagen-4-fast",
"imagen-4-ultra",
"dall-e-3",
"gpt-image-1",
]
for model_id in model_ids:
info = get_model_info(model_id)
assert info is not None
assert info["id"] == model_id
class TestListModelIds:
"""Tests for list_model_ids function."""
def test_list_model_ids_returns_list(self):
"""list_model_ids should return a list of strings."""
ids = list_model_ids()
assert isinstance(ids, list)
assert all(isinstance(id, str) for id in ids)
def test_list_model_ids_contains_all_models(self):
"""list_model_ids should contain all model IDs."""
ids = list_model_ids()
assert len(ids) == 5
assert "imagen-4" in ids
assert "imagen-4-fast" in ids
assert "imagen-4-ultra" in ids
assert "dall-e-3" in ids
assert "gpt-image-1" in ids
