"""Unit tests for Nexos.ai client."""
import pytest
from pytest_httpx import HTTPXMock
from Imagen_MCP.services.nexos_client import NexosClient
from Imagen_MCP.models.generation import GenerateImageRequest
from Imagen_MCP.exceptions import (
AuthenticationError,
ConfigurationError,
GenerationError,
InvalidRequestError,
RateLimitError,
)
from tests.mocks.nexos_api import (
MOCK_IMAGE_GENERATION_SUCCESS,
MOCK_IMAGE_GENERATION_RATE_LIMITED,
MOCK_MODELS_LIST,
)
class TestNexosClientInitialization:
"""Tests for client initialization."""
def test_client_initialization_with_api_key(self):
"""Client should initialize with valid API key."""
client = NexosClient(api_key="test-api-key")
assert client.api_key == "test-api-key"
assert client._base_url == "https://api.nexos.ai/v1"
def test_client_initialization_with_custom_base_url(self):
"""Client should accept custom base URL."""
client = NexosClient(
api_key="test-api-key",
base_url="https://custom.api.com/v2/",
)
assert client._base_url == "https://custom.api.com/v2"
def test_client_initialization_without_api_key_raises(self):
"""Client should raise ConfigurationError when API key is missing."""
client = NexosClient()
with pytest.raises(ConfigurationError, match="NEXOS_API_KEY is not configured"):
_ = client.api_key
class TestGenerateImage:
"""Tests for image generation."""
@pytest.fixture
def client(self):
"""Create a test client."""
return NexosClient(api_key="test-api-key")
@pytest.fixture
def sample_request(self):
"""Create a sample generation request."""
return GenerateImageRequest(
prompt="A serene mountain landscape at sunset",
model="imagen-4",
size="1024x1024",
quality="standard",
style="vivid",
)
@pytest.mark.asyncio
async def test_generate_image_request_format(
self,
client: NexosClient,
sample_request: GenerateImageRequest,
httpx_mock: HTTPXMock,
):
"""Generate image should send correctly formatted request to API."""
httpx_mock.add_response(
url="https://api.nexos.ai/v1/images/generations",
method="POST",
json=MOCK_IMAGE_GENERATION_SUCCESS,
)
await client.generate_image(sample_request)
# Verify the request was made correctly
requests = httpx_mock.get_requests()
assert len(requests) == 1
request = requests[0]
assert request.method == "POST"
assert "Bearer test-api-key" in request.headers["authorization"]
@pytest.mark.asyncio
async def test_generate_image_response_parsing(
self,
client: NexosClient,
sample_request: GenerateImageRequest,
httpx_mock: HTTPXMock,
):
"""Client should correctly parse API response into GeneratedImage."""
httpx_mock.add_response(
url="https://api.nexos.ai/v1/images/generations",
method="POST",
json=MOCK_IMAGE_GENERATION_SUCCESS,
)
response = await client.generate_image(sample_request)
assert response.created == 1734800000
assert len(response.images) == 1
assert response.images[0].b64_json is not None
assert response.images[0].revised_prompt is not None
@pytest.mark.asyncio
async def test_generate_image_with_all_parameters(
self, client: NexosClient, httpx_mock: HTTPXMock
):
"""All optional parameters should be included in API request."""
httpx_mock.add_response(
url="https://api.nexos.ai/v1/images/generations",
method="POST",
json=MOCK_IMAGE_GENERATION_SUCCESS,
)
request = GenerateImageRequest(
prompt="Test prompt",
model="imagen-4-ultra",
size="1792x1024",
quality="hd",
style="natural",
n=2,
response_format="b64_json",
)
await client.generate_image(request)
# Verify request payload
requests = httpx_mock.get_requests()
import json
payload = json.loads(requests[0].content)
assert payload["prompt"] == "Test prompt"
assert payload["model"] == "imagen-4-ultra"
assert payload["size"] == "1792x1024"
assert payload["quality"] == "hd"
assert payload["style"] == "natural"
assert payload["n"] == 2
class TestErrorHandling:
"""Tests for API error handling."""
@pytest.fixture
def client(self):
"""Create a test client."""
return NexosClient(api_key="test-api-key")
@pytest.fixture
def sample_request(self):
"""Create a sample generation request."""
return GenerateImageRequest(prompt="Test prompt")
@pytest.mark.asyncio
async def test_api_error_400_handling(
self,
client: NexosClient,
sample_request: GenerateImageRequest,
httpx_mock: HTTPXMock,
):
"""Client should raise InvalidRequestError on 400 response."""
httpx_mock.add_response(
url="https://api.nexos.ai/v1/images/generations",
method="POST",
status_code=400,
json={
"error": {"message": "Invalid prompt", "type": "invalid_request_error"}
},
)
with pytest.raises(InvalidRequestError, match="Invalid request"):
await client.generate_image(sample_request)
@pytest.mark.asyncio
async def test_api_error_401_handling(
self,
client: NexosClient,
sample_request: GenerateImageRequest,
httpx_mock: HTTPXMock,
):
"""Client should raise AuthenticationError on 401 response."""
httpx_mock.add_response(
url="https://api.nexos.ai/v1/images/generations",
method="POST",
status_code=401,
json={
"error": {"message": "Invalid API key", "type": "authentication_error"}
},
)
with pytest.raises(AuthenticationError, match="Authentication failed"):
await client.generate_image(sample_request)
@pytest.mark.asyncio
async def test_api_error_429_handling(
self,
client: NexosClient,
sample_request: GenerateImageRequest,
httpx_mock: HTTPXMock,
):
"""Client should raise RateLimitError on 429 response."""
httpx_mock.add_response(
url="https://api.nexos.ai/v1/images/generations",
method="POST",
status_code=429,
json=MOCK_IMAGE_GENERATION_RATE_LIMITED,
headers={"retry-after": "60"},
)
with pytest.raises(RateLimitError, match="Rate limit exceeded") as exc_info:
await client.generate_image(sample_request)
assert exc_info.value.retry_after == 60
@pytest.mark.asyncio
async def test_api_error_500_handling(
self,
client: NexosClient,
sample_request: GenerateImageRequest,
httpx_mock: HTTPXMock,
):
"""Client should raise GenerationError on 500 response."""
# Add multiple responses for retry attempts
for _ in range(4): # MAX_RETRIES + 1
httpx_mock.add_response(
url="https://api.nexos.ai/v1/images/generations",
method="POST",
status_code=500,
json={
"error": {
"message": "Internal server error",
"type": "server_error",
}
},
)
with pytest.raises(GenerationError, match="Server error"):
await client.generate_image(sample_request)
class TestRetryLogic:
"""Tests for retry logic on transient errors."""
@pytest.fixture
def client(self):
"""Create a test client with minimal retry delay."""
client = NexosClient(api_key="test-api-key")
client.RETRY_DELAY = 0.01 # Speed up tests
return client
@pytest.fixture
def sample_request(self):
"""Create a sample generation request."""
return GenerateImageRequest(prompt="Test prompt")
@pytest.mark.asyncio
async def test_retry_on_transient_errors(
self,
client: NexosClient,
sample_request: GenerateImageRequest,
httpx_mock: HTTPXMock,
):
"""Client should retry on 5xx errors with exponential backoff."""
# First two requests fail, third succeeds
httpx_mock.add_response(
url="https://api.nexos.ai/v1/images/generations",
method="POST",
status_code=500,
json={"error": {"message": "Server error", "type": "server_error"}},
)
httpx_mock.add_response(
url="https://api.nexos.ai/v1/images/generations",
method="POST",
status_code=500,
json={"error": {"message": "Server error", "type": "server_error"}},
)
httpx_mock.add_response(
url="https://api.nexos.ai/v1/images/generations",
method="POST",
json=MOCK_IMAGE_GENERATION_SUCCESS,
)
response = await client.generate_image(sample_request)
# Should have made 3 requests
requests = httpx_mock.get_requests()
assert len(requests) == 3
assert len(response.images) == 1
class TestListModels:
"""Tests for listing models."""
@pytest.fixture
def client(self):
"""Create a test client."""
return NexosClient(api_key="test-api-key")
@pytest.mark.asyncio
async def test_list_models(self, client: NexosClient, httpx_mock: HTTPXMock):
"""Client should correctly fetch and parse available models."""
httpx_mock.add_response(
url="https://api.nexos.ai/v1/models",
method="GET",
json=MOCK_MODELS_LIST,
)
models = await client.list_models()
assert len(models) == 4
assert models[0]["id"] == "imagen-4"
assert models[1]["id"] == "imagen-4-fast"
