"""Shared pytest fixtures for Imagen MCP tests."""
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock
import pytest
from Imagen_MCP.services.session_manager import SessionManager
from Imagen_MCP.services.nexos_client import NexosClient
from Imagen_MCP.models.generation import GenerateImageResponse
from Imagen_MCP.models.image import GeneratedImage
# Test image data - 1x1 pixel transparent PNG
TEST_B64_PNG = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" # pragma: allowlist secret
@pytest.fixture
def test_b64_png() -> str:
"""Provide test base64-encoded PNG data."""
return TEST_B64_PNG
@pytest.fixture
def temp_output_dir(tmp_path: Path) -> Path:
"""Provide a temporary directory for test output."""
output_dir = tmp_path / "imagen_test_output"
output_dir.mkdir(parents=True, exist_ok=True)
return output_dir
@pytest.fixture
def temp_output_file(temp_output_dir: Path) -> Path:
"""Provide a temporary file path for test output."""
return temp_output_dir / "test_image.png"
@pytest.fixture
def sample_generation_request() -> dict:
"""Provide a sample generation request for testing."""
return {
"prompt": "A serene mountain landscape at sunset",
"model": "imagen-4",
"size": "1024x1024",
"quality": "standard",
"style": "vivid",
}
@pytest.fixture
def mock_generated_image(test_b64_png: str) -> GeneratedImage:
"""Provide a mock GeneratedImage for testing."""
return GeneratedImage(
b64_json=test_b64_png,
revised_prompt="A serene mountain landscape at sunset with golden light",
model="imagen-4",
size="1024x1024",
quality="standard",
style="vivid",
)
@pytest.fixture
def mock_generation_response(
mock_generated_image: GeneratedImage,
) -> GenerateImageResponse:
"""Provide a mock GenerateImageResponse for testing."""
return GenerateImageResponse(
created=1734800000,
images=[mock_generated_image],
)
@pytest.fixture
def mock_nexos_client(mock_generation_response: GenerateImageResponse) -> MagicMock:
"""Provide a mock NexosClient for testing."""
client = MagicMock(spec=NexosClient)
client.generate_image = AsyncMock(return_value=mock_generation_response)
client.__aenter__ = AsyncMock(return_value=client)
client.__aexit__ = AsyncMock(return_value=None)
return client
@pytest.fixture
def mock_session_manager() -> SessionManager:
"""Provide a fresh SessionManager instance for testing."""
return SessionManager(
client=None,
max_concurrent_sessions=10,
session_ttl_minutes=60,
)
@pytest.fixture
def api_success_response(test_b64_png: str) -> dict:
"""Provide a mock successful API response."""
return {
"created": 1734800000,
"data": [
{
"b64_json": test_b64_png,
"revised_prompt": "A serene mountain landscape at sunset with golden light",
}
],
}
@pytest.fixture
def api_error_response() -> dict:
"""Provide a mock error API response."""
return {
"error": {
"message": "Rate limit exceeded",
"type": "rate_limit_error",
"code": "rate_limit_exceeded",
}
}
