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()
584 lines
18 KiB
Python
584 lines
18 KiB
Python
"""
|
|
Unit tests for data models (entities and schemas)
|
|
|
|
Tests entity creation, validation, and schema validation rules.
|
|
"""
|
|
import pytest
|
|
from datetime import datetime
|
|
import uuid
|
|
from typing import Dict, Any
|
|
|
|
from src.models.entities import (
|
|
ProjectDB,
|
|
AssetDB,
|
|
EpisodeDB,
|
|
StoryboardDB,
|
|
TaskDB,
|
|
CanvasDB,
|
|
CanvasMetadataDB,
|
|
)
|
|
from src.models.schemas import (
|
|
CreateProjectRequest,
|
|
UpdateProjectRequest,
|
|
CreateCharacterAssetRequest,
|
|
CreateSceneAssetRequest,
|
|
CreatePropAssetRequest,
|
|
CreateEpisodeRequest,
|
|
UpdateEpisodeRequest,
|
|
CreateStoryboardRequest,
|
|
UpdateStoryboardRequest,
|
|
ImageGenerationRequest,
|
|
VideoGenerationRequest,
|
|
Task,
|
|
CanvasMetadata,
|
|
)
|
|
from pydantic import ValidationError
|
|
|
|
|
|
class TestProjectDBEntity:
|
|
"""Test ProjectDB entity creation and validation"""
|
|
|
|
def test_create_project_with_defaults(self):
|
|
"""Test creating a project with default values"""
|
|
project = ProjectDB(
|
|
name="Test Project",
|
|
description="Test Description"
|
|
)
|
|
|
|
assert project.name == "Test Project"
|
|
assert project.description == "Test Description"
|
|
assert project.type == "video"
|
|
assert project.status == "active"
|
|
assert project.id is not None
|
|
assert project.created_at is not None
|
|
assert project.updated_at is not None
|
|
assert project.deleted_at is None
|
|
|
|
def test_create_project_with_custom_values(self):
|
|
"""Test creating a project with custom values"""
|
|
custom_id = str(uuid.uuid4())
|
|
custom_time = datetime.now().timestamp()
|
|
|
|
project = ProjectDB(
|
|
id=custom_id,
|
|
name="Custom Project",
|
|
type="video",
|
|
status="initializing",
|
|
created_at=custom_time,
|
|
updated_at=custom_time,
|
|
resolution="1920x1080",
|
|
ratio="16:9",
|
|
style_id="anime",
|
|
style_params={"color": "vibrant"}
|
|
)
|
|
|
|
assert project.id == custom_id
|
|
assert project.name == "Custom Project"
|
|
assert project.type == "video"
|
|
assert project.status == "initializing"
|
|
assert project.resolution == "1920x1080"
|
|
assert project.ratio == "16:9"
|
|
assert project.style_id == "anime"
|
|
assert project.style_params == {"color": "vibrant"}
|
|
|
|
def test_project_with_json_fields(self):
|
|
"""Test project with JSON fields (chapters, progress, error)"""
|
|
chapters = [
|
|
{"title": "Chapter 1", "content": "Content 1"},
|
|
{"title": "Chapter 2", "content": "Content 2"}
|
|
]
|
|
progress = {"step": "analyzing", "percentage": 50}
|
|
error = {"code": "E001", "message": "Test error"}
|
|
|
|
project = ProjectDB(
|
|
name="Test Project",
|
|
chapters=chapters,
|
|
progress=progress,
|
|
error=error
|
|
)
|
|
|
|
assert project.chapters == chapters
|
|
assert project.progress == progress
|
|
assert project.error == error
|
|
|
|
|
|
class TestAssetDBEntity:
|
|
"""Test AssetDB entity creation and validation"""
|
|
|
|
def test_create_asset_with_required_fields(self):
|
|
"""Test creating an asset with required fields only"""
|
|
asset = AssetDB(
|
|
project_id="proj_123",
|
|
type="character",
|
|
name="Hero"
|
|
)
|
|
|
|
assert asset.id is not None
|
|
assert asset.project_id == "proj_123"
|
|
assert asset.type == "character"
|
|
assert asset.name == "Hero"
|
|
assert asset.desc == ""
|
|
assert asset.tags == []
|
|
assert asset.extra_data == {}
|
|
assert asset.generations == []
|
|
|
|
def test_create_asset_with_all_fields(self):
|
|
"""Test creating an asset with all fields"""
|
|
asset = AssetDB(
|
|
project_id="proj_123",
|
|
type="character",
|
|
name="Hero",
|
|
desc="Main character",
|
|
tags=["protagonist", "hero"],
|
|
image_url="https://example.com/hero.png",
|
|
image_urls=["https://example.com/hero1.png", "https://example.com/hero2.png"],
|
|
video_urls=["https://example.com/hero.mp4"],
|
|
image_prompt="A heroic character",
|
|
extra_data={"age": "25", "gender": "male"}
|
|
)
|
|
|
|
assert asset.name == "Hero"
|
|
assert asset.desc == "Main character"
|
|
assert len(asset.tags) == 2
|
|
assert asset.image_url == "https://example.com/hero.png"
|
|
assert len(asset.image_urls) == 2
|
|
assert asset.extra_data["age"] == "25"
|
|
|
|
|
|
class TestEpisodeDBEntity:
|
|
"""Test EpisodeDB entity creation and validation"""
|
|
|
|
def test_create_episode(self):
|
|
"""Test creating an episode"""
|
|
episode = EpisodeDB(
|
|
project_id="proj_123",
|
|
order_index=1,
|
|
title="Episode 1",
|
|
desc="First episode",
|
|
content="Episode content",
|
|
status="draft"
|
|
)
|
|
|
|
assert episode.id is not None
|
|
assert episode.project_id == "proj_123"
|
|
assert episode.order_index == 1
|
|
assert episode.title == "Episode 1"
|
|
assert episode.desc == "First episode"
|
|
assert episode.status == "draft"
|
|
|
|
|
|
class TestStoryboardDBEntity:
|
|
"""Test StoryboardDB entity creation and validation"""
|
|
|
|
def test_create_storyboard_minimal(self):
|
|
"""Test creating a storyboard with minimal fields"""
|
|
storyboard = StoryboardDB(
|
|
project_id="proj_123",
|
|
episode_id="ep_123",
|
|
order_index=1,
|
|
shot="Shot 1",
|
|
desc="Opening scene",
|
|
duration="5s",
|
|
type="image"
|
|
)
|
|
|
|
assert storyboard.id is not None
|
|
assert storyboard.project_id == "proj_123"
|
|
assert storyboard.episode_id == "ep_123"
|
|
assert storyboard.shot == "Shot 1"
|
|
assert storyboard.character_ids == []
|
|
assert storyboard.prop_ids == []
|
|
|
|
def test_create_storyboard_with_cinematic_fields(self):
|
|
"""Test creating a storyboard with cinematic control fields"""
|
|
storyboard = StoryboardDB(
|
|
project_id="proj_123",
|
|
episode_id="ep_123",
|
|
order_index=1,
|
|
shot="Shot 1",
|
|
desc="Opening scene",
|
|
duration="5s",
|
|
type="image",
|
|
camera_angle="wide",
|
|
lens="50mm",
|
|
focus="deep",
|
|
lighting="natural",
|
|
color_style="warm",
|
|
location="forest",
|
|
time="morning"
|
|
)
|
|
|
|
assert storyboard.camera_angle == "wide"
|
|
assert storyboard.lens == "50mm"
|
|
assert storyboard.focus == "deep"
|
|
assert storyboard.lighting == "natural"
|
|
assert storyboard.color_style == "warm"
|
|
assert storyboard.location == "forest"
|
|
assert storyboard.time == "morning"
|
|
|
|
|
|
class TestTaskDBEntity:
|
|
"""Test TaskDB entity creation and validation"""
|
|
|
|
def test_create_task_with_defaults(self):
|
|
"""Test creating a task with default values"""
|
|
task = TaskDB(
|
|
type="image",
|
|
status="pending",
|
|
model="flux-dev",
|
|
params={"prompt": "test"}
|
|
)
|
|
|
|
assert task.id is not None
|
|
assert task.type == "image"
|
|
assert task.status == "pending"
|
|
assert task.retry_count == 0
|
|
assert task.max_retries == 3
|
|
assert task.created_at is not None
|
|
assert task.updated_at is not None
|
|
assert task.started_at is None
|
|
assert task.completed_at is None
|
|
|
|
def test_create_task_with_all_fields(self):
|
|
"""Test creating a task with all fields"""
|
|
now = datetime.now().timestamp()
|
|
task = TaskDB(
|
|
type="video",
|
|
status="processing",
|
|
model="kling-v1",
|
|
params={"prompt": "test video"},
|
|
provider_task_id="provider_123",
|
|
result={"url": "https://example.com/video.mp4"},
|
|
error=None,
|
|
retry_count=1,
|
|
max_retries=5,
|
|
started_at=now,
|
|
user_id="user_123",
|
|
project_id="proj_123"
|
|
)
|
|
|
|
assert task.type == "video"
|
|
assert task.status == "processing"
|
|
assert task.provider_task_id == "provider_123"
|
|
assert task.result["url"] == "https://example.com/video.mp4"
|
|
assert task.retry_count == 1
|
|
assert task.max_retries == 5
|
|
assert task.user_id == "user_123"
|
|
|
|
|
|
class TestCanvasMetadataDBEntity:
|
|
"""Test CanvasMetadataDB entity creation and validation"""
|
|
|
|
def test_create_general_canvas_metadata(self):
|
|
"""Test creating general canvas metadata"""
|
|
canvas = CanvasMetadataDB(
|
|
project_id="proj_123",
|
|
canvas_type="general",
|
|
name="Main Canvas",
|
|
description="Main project canvas",
|
|
order_index=0,
|
|
is_pinned=True,
|
|
tags=["main", "primary"]
|
|
)
|
|
|
|
assert canvas.id is not None
|
|
assert canvas.project_id == "proj_123"
|
|
assert canvas.canvas_type == "general"
|
|
assert canvas.name == "Main Canvas"
|
|
assert canvas.is_pinned is True
|
|
assert len(canvas.tags) == 2
|
|
assert canvas.node_count == 0
|
|
assert canvas.access_count == 0
|
|
|
|
def test_create_asset_canvas_metadata(self):
|
|
"""Test creating asset-related canvas metadata"""
|
|
canvas = CanvasMetadataDB(
|
|
project_id="proj_123",
|
|
canvas_type="asset",
|
|
related_entity_type="asset",
|
|
related_entity_id="asset_123",
|
|
name="Character Canvas",
|
|
order_index=1,
|
|
is_pinned=False
|
|
)
|
|
|
|
assert canvas.canvas_type == "asset"
|
|
assert canvas.related_entity_type == "asset"
|
|
assert canvas.related_entity_id == "asset_123"
|
|
|
|
|
|
class TestProjectSchemas:
|
|
"""Test Project schema validation"""
|
|
|
|
def test_create_project_request_valid(self):
|
|
"""Test valid CreateProjectRequest"""
|
|
request = CreateProjectRequest(
|
|
name="Test Project",
|
|
description="Test Description",
|
|
type="video"
|
|
)
|
|
|
|
assert request.name == "Test Project"
|
|
assert request.description == "Test Description"
|
|
assert request.type == "video"
|
|
|
|
def test_create_project_request_minimal(self):
|
|
"""Test CreateProjectRequest with minimal fields"""
|
|
request = CreateProjectRequest(name="Test Project")
|
|
|
|
assert request.name == "Test Project"
|
|
assert request.description is None
|
|
assert request.type == "video" # default value
|
|
|
|
def test_create_project_request_invalid_type(self):
|
|
"""Test CreateProjectRequest accepts any type value (no strict validation)"""
|
|
# Note: type field doesn't have strict validation, so any string is accepted
|
|
request = CreateProjectRequest(
|
|
name="Test Project",
|
|
type="custom_type"
|
|
)
|
|
|
|
assert request.type == "custom_type"
|
|
|
|
def test_update_project_request(self):
|
|
"""Test UpdateProjectRequest"""
|
|
request = UpdateProjectRequest(
|
|
name="Updated Name",
|
|
resolution="1920x1080",
|
|
styleId="anime"
|
|
)
|
|
|
|
assert request.name == "Updated Name"
|
|
assert request.resolution == "1920x1080"
|
|
assert request.style_id == "anime"
|
|
assert request.description is None # not updated
|
|
|
|
|
|
class TestAssetSchemas:
|
|
"""Test Asset schema validation"""
|
|
|
|
def test_create_character_asset_request(self):
|
|
"""Test CreateCharacterAssetRequest"""
|
|
request = CreateCharacterAssetRequest(
|
|
type="character",
|
|
name="Hero",
|
|
desc="Main character",
|
|
tags=["protagonist"],
|
|
age="25",
|
|
gender="male",
|
|
role="hero"
|
|
)
|
|
|
|
assert request.type == "character"
|
|
assert request.name == "Hero"
|
|
assert request.age == "25"
|
|
assert request.gender == "male"
|
|
|
|
def test_create_scene_asset_request(self):
|
|
"""Test CreateSceneAssetRequest"""
|
|
request = CreateSceneAssetRequest(
|
|
type="scene",
|
|
name="Forest",
|
|
desc="Dark forest",
|
|
location="Northern Woods",
|
|
time_of_day="night",
|
|
atmosphere="mysterious"
|
|
)
|
|
|
|
assert request.type == "scene"
|
|
assert request.name == "Forest"
|
|
assert request.location == "Northern Woods"
|
|
assert request.time_of_day == "night"
|
|
|
|
def test_create_prop_asset_request(self):
|
|
"""Test CreatePropAssetRequest"""
|
|
request = CreatePropAssetRequest(
|
|
type="prop",
|
|
name="Magic Sword",
|
|
desc="Ancient sword",
|
|
usage="weapon"
|
|
)
|
|
|
|
assert request.type == "prop"
|
|
assert request.name == "Magic Sword"
|
|
assert request.usage == "weapon"
|
|
|
|
|
|
class TestEpisodeSchemas:
|
|
"""Test Episode schema validation"""
|
|
|
|
def test_create_episode_request(self):
|
|
"""Test CreateEpisodeRequest"""
|
|
request = CreateEpisodeRequest(
|
|
title="Episode 1",
|
|
order=1,
|
|
desc="First episode",
|
|
status="draft"
|
|
)
|
|
|
|
assert request.title == "Episode 1"
|
|
assert request.order == 1
|
|
assert request.status == "draft"
|
|
|
|
def test_update_episode_request(self):
|
|
"""Test UpdateEpisodeRequest"""
|
|
request = UpdateEpisodeRequest(
|
|
title="Updated Episode",
|
|
status="production"
|
|
)
|
|
|
|
assert request.title == "Updated Episode"
|
|
assert request.status == "production"
|
|
assert request.order is None # not updated
|
|
|
|
|
|
class TestStoryboardSchemas:
|
|
"""Test Storyboard schema validation"""
|
|
|
|
def test_create_storyboard_request(self):
|
|
"""Test CreateStoryboardRequest"""
|
|
request = CreateStoryboardRequest(
|
|
episode_id="ep_123",
|
|
order=1,
|
|
shot="Shot 1",
|
|
desc="Opening scene",
|
|
duration="5s",
|
|
type="image",
|
|
scene_id="scene_123",
|
|
character_ids=["char_1", "char_2"],
|
|
camera_angle="wide"
|
|
)
|
|
|
|
assert request.episode_id == "ep_123"
|
|
assert request.order == 1
|
|
assert request.shot == "Shot 1"
|
|
assert len(request.character_ids) == 2
|
|
assert request.camera_angle == "wide"
|
|
|
|
def test_update_storyboard_request(self):
|
|
"""Test UpdateStoryboardRequest"""
|
|
request = UpdateStoryboardRequest(
|
|
shot="Updated Shot",
|
|
duration="10s",
|
|
camera_angle="close-up"
|
|
)
|
|
|
|
assert request.shot == "Updated Shot"
|
|
assert request.duration == "10s"
|
|
assert request.camera_angle == "close-up"
|
|
|
|
|
|
class TestGenerationSchemas:
|
|
"""Test Generation request schema validation"""
|
|
|
|
def test_image_generation_request_minimal(self):
|
|
"""Test ImageGenerationRequest with minimal fields"""
|
|
request = ImageGenerationRequest(
|
|
prompt="A beautiful landscape",
|
|
model="replicate/flux-dev" # Format: provider/model_key
|
|
)
|
|
|
|
assert request.prompt == "A beautiful landscape"
|
|
assert request.n == 1 # default
|
|
assert request.model == "replicate/flux-dev"
|
|
|
|
def test_image_generation_request_full(self):
|
|
"""Test ImageGenerationRequest with all fields"""
|
|
request = ImageGenerationRequest(
|
|
prompt="A beautiful landscape",
|
|
negativePrompt="ugly, blurry",
|
|
model="dashscope/flux-dev",
|
|
imageInputs=["https://example.com/ref.png"],
|
|
resolution="1K", # Quality level format
|
|
aspectRatio="1:1",
|
|
n=2,
|
|
projectId="proj_123"
|
|
)
|
|
|
|
assert request.prompt == "A beautiful landscape"
|
|
assert request.negative_prompt == "ugly, blurry"
|
|
assert request.model == "dashscope/flux-dev"
|
|
assert request.n == 2
|
|
assert request.image_inputs == ["https://example.com/ref.png"]
|
|
assert request.resolution == "1K"
|
|
|
|
def test_video_generation_request_minimal(self):
|
|
"""Test VideoGenerationRequest with minimal fields"""
|
|
request = VideoGenerationRequest(
|
|
prompt="A flowing river",
|
|
model="kling/v1" # Format: provider/model_key
|
|
)
|
|
|
|
assert request.prompt == "A flowing river"
|
|
assert request.duration == 5 # default
|
|
|
|
def test_video_generation_request_with_images(self):
|
|
"""Test VideoGenerationRequest with image URLs"""
|
|
request = VideoGenerationRequest(
|
|
imageInputs=["https://example.com/img1.png", "https://example.com/img2.png"],
|
|
duration=10,
|
|
aspectRatio="16:9",
|
|
model="kling/v1",
|
|
prompt="A flowing river"
|
|
)
|
|
|
|
assert request.image_inputs == ["https://example.com/img1.png", "https://example.com/img2.png"]
|
|
assert request.duration == 10
|
|
assert request.aspect_ratio == "16:9"
|
|
|
|
|
|
class TestTaskSchema:
|
|
"""Test Task schema validation"""
|
|
|
|
def test_task_schema(self):
|
|
"""Test Task schema"""
|
|
now = datetime.now().timestamp()
|
|
task = Task(
|
|
id="task_123",
|
|
type="image",
|
|
status="pending",
|
|
created_at=now,
|
|
updated_at=now,
|
|
model="flux-dev",
|
|
params={"prompt": "test"},
|
|
retry_count=0,
|
|
max_retries=3
|
|
)
|
|
|
|
assert task.id == "task_123"
|
|
assert task.type == "image"
|
|
assert task.status == "pending"
|
|
assert task.model == "flux-dev"
|
|
assert task.retry_count == 0
|
|
|
|
|
|
class TestCanvasMetadataSchema:
|
|
"""Test CanvasMetadata schema validation"""
|
|
|
|
def test_canvas_metadata_schema(self):
|
|
"""Test CanvasMetadata schema with alias fields"""
|
|
now = datetime.now().timestamp()
|
|
canvas = CanvasMetadata(
|
|
id="canvas_123",
|
|
projectId="proj_123",
|
|
canvasType="general",
|
|
name="Main Canvas",
|
|
orderIndex=0,
|
|
isPinned=True,
|
|
tags=["main"],
|
|
nodeCount=5,
|
|
accessCount=10,
|
|
createdAt=now,
|
|
updatedAt=now
|
|
)
|
|
|
|
assert canvas.id == "canvas_123"
|
|
assert canvas.project_id == "proj_123"
|
|
assert canvas.canvas_type == "general"
|
|
assert canvas.is_pinned is True
|
|
assert canvas.node_count == 5
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|