f9f4560459
- FastAPI backend with SQLModel, Alembic migrations, AgentScope agents - Next.js 15 frontend with React 19, Tailwind, Zustand, React Flow - Multi-provider AI system (DashScope, Kling, MiniMax, Volcengine, OpenAI, etc.) - All HTTP clients migrated from sync requests to async httpx - Admin-managed API keys via environment variables - SSRF vulnerability fixed in ensure_url()
149 lines
4.9 KiB
Python
149 lines
4.9 KiB
Python
"""
|
|
Tests for ImageGenerationRequest schema validation.
|
|
|
|
Tests the model field format validation to ensure it accepts valid
|
|
composite IDs (provider/model_key) and rejects invalid formats.
|
|
"""
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
from src.models.schemas import ImageGenerationRequest
|
|
from src.utils.errors import InvalidParameterException
|
|
|
|
|
|
class TestImageGenerationRequestModelValidation:
|
|
"""Test model field format validation"""
|
|
|
|
def test_accepts_valid_composite_id(self):
|
|
"""Should accept valid composite ID format"""
|
|
request = ImageGenerationRequest(
|
|
prompt="a cat",
|
|
model="dashscope/qwen-image"
|
|
)
|
|
assert request.model == "dashscope/qwen-image"
|
|
|
|
def test_accepts_different_providers(self):
|
|
"""Should accept different provider formats"""
|
|
valid_models = [
|
|
"dashscope/qwen-image",
|
|
"modelscope/qwen-image",
|
|
"volcengine/doubao-image",
|
|
"openai/dall-e-3"
|
|
]
|
|
|
|
for model in valid_models:
|
|
request = ImageGenerationRequest(
|
|
prompt="test",
|
|
model=model
|
|
)
|
|
assert request.model == model
|
|
|
|
def test_rejects_model_without_separator(self):
|
|
"""Should reject model without '/' separator"""
|
|
with pytest.raises((ValidationError, InvalidParameterException)):
|
|
ImageGenerationRequest(
|
|
prompt="a cat",
|
|
model="qwen-image"
|
|
)
|
|
|
|
def test_rejects_model_with_multiple_separators(self):
|
|
"""Should reject model with multiple '/' separators"""
|
|
with pytest.raises((ValidationError, InvalidParameterException)):
|
|
ImageGenerationRequest(
|
|
prompt="a cat",
|
|
model="dash/scope/qwen"
|
|
)
|
|
|
|
def test_rejects_empty_provider(self):
|
|
"""Should reject model with empty provider"""
|
|
with pytest.raises((ValidationError, InvalidParameterException)):
|
|
ImageGenerationRequest(
|
|
prompt="a cat",
|
|
model="/qwen-image"
|
|
)
|
|
|
|
def test_rejects_empty_model_key(self):
|
|
"""Should reject model with empty model_key"""
|
|
with pytest.raises((ValidationError, InvalidParameterException)):
|
|
ImageGenerationRequest(
|
|
prompt="a cat",
|
|
model="dashscope/"
|
|
)
|
|
|
|
def test_provider_field_removed(self):
|
|
"""Should not accept provider field (removed)"""
|
|
# This should work without provider field
|
|
request = ImageGenerationRequest(
|
|
prompt="a cat",
|
|
model="dashscope/qwen-image"
|
|
)
|
|
assert not hasattr(request, 'provider') or request.model_dump().get('provider') is None
|
|
|
|
def test_complete_request_with_optional_fields(self):
|
|
"""Should accept complete request with all optional fields"""
|
|
request = ImageGenerationRequest(
|
|
prompt="a beautiful cat",
|
|
model="dashscope/qwen-image",
|
|
negative_prompt="ugly",
|
|
image_inputs=["http://example.com/image.jpg"],
|
|
resolution="2K",
|
|
aspect_ratio="16:9",
|
|
n=2,
|
|
project_id="proj-123",
|
|
source="storyboard",
|
|
source_id="story-456",
|
|
extra_params={"lora": "style1"}
|
|
)
|
|
|
|
assert request.model == "dashscope/qwen-image"
|
|
assert request.prompt == "a beautiful cat"
|
|
assert request.n == 2
|
|
|
|
|
|
class TestImageGenerationRequestOtherValidations:
|
|
"""Test other field validations"""
|
|
|
|
def test_prompt_required(self):
|
|
"""Should require prompt field"""
|
|
with pytest.raises(ValidationError) as exc_info:
|
|
ImageGenerationRequest(
|
|
model="dashscope/qwen-image"
|
|
)
|
|
|
|
errors = exc_info.value.errors()
|
|
assert any(e["loc"] == ("prompt",) for e in errors)
|
|
|
|
def test_prompt_cannot_be_empty(self):
|
|
"""Should reject empty prompt"""
|
|
with pytest.raises((ValidationError, InvalidParameterException)):
|
|
ImageGenerationRequest(
|
|
prompt="",
|
|
model="dashscope/qwen-image"
|
|
)
|
|
|
|
def test_n_validation(self):
|
|
"""Should validate n is between 1 and 10"""
|
|
# Valid n
|
|
request = ImageGenerationRequest(
|
|
prompt="test",
|
|
model="dashscope/qwen-image",
|
|
n=5
|
|
)
|
|
assert request.n == 5
|
|
|
|
# Invalid n (too small)
|
|
with pytest.raises((ValidationError, InvalidParameterException)):
|
|
ImageGenerationRequest(
|
|
prompt="test",
|
|
model="dashscope/qwen-image",
|
|
n=0
|
|
)
|
|
|
|
# Invalid n (too large)
|
|
with pytest.raises((ValidationError, InvalidParameterException)):
|
|
ImageGenerationRequest(
|
|
prompt="test",
|
|
model="dashscope/qwen-image",
|
|
n=11
|
|
)
|