180 lines
6.9 KiB
Python
180 lines
6.9 KiB
Python
"""
|
|
Test custom task type integration with asset and shot creation
|
|
"""
|
|
import sys
|
|
import os
|
|
sys.path.insert(0, os.path.dirname(__file__))
|
|
|
|
from database import SessionLocal, engine
|
|
from models.project import Project
|
|
from models.asset import Asset, AssetCategory
|
|
from models.episode import Episode
|
|
from models.shot import Shot
|
|
from models.task import Task
|
|
from models.user import User, UserRole
|
|
from sqlalchemy.orm.attributes import flag_modified
|
|
import json
|
|
|
|
def test_custom_task_types_integration():
|
|
"""Test that custom task types work with asset and shot creation"""
|
|
db = SessionLocal()
|
|
|
|
try:
|
|
print("\n=== Testing Custom Task Types Integration ===\n")
|
|
|
|
# 1. Get or create a test project with custom task types
|
|
print("1. Getting/creating test project with custom task types...")
|
|
project = db.query(Project).first()
|
|
|
|
if not project:
|
|
print(" No existing project found, creating new one...")
|
|
project = Project(
|
|
name="Custom Task Test Project",
|
|
status="planning"
|
|
)
|
|
db.add(project)
|
|
db.commit()
|
|
db.refresh(project)
|
|
|
|
# Add custom task types to the project
|
|
project.custom_asset_task_types = ["grooming", "lookdev"]
|
|
project.custom_shot_task_types = ["previz", "matchmove"]
|
|
flag_modified(project, 'custom_asset_task_types')
|
|
flag_modified(project, 'custom_shot_task_types')
|
|
db.commit()
|
|
db.refresh(project)
|
|
|
|
print(f" ✓ Using project: {project.name} (ID: {project.id})")
|
|
print(f" ✓ Custom asset types: {project.custom_asset_task_types}")
|
|
print(f" ✓ Custom shot types: {project.custom_shot_task_types}")
|
|
|
|
# 2. Create an asset with custom task type
|
|
print("\n2. Creating asset with custom task type...")
|
|
asset = Asset(
|
|
project_id=project.id,
|
|
name="Test Character",
|
|
category=AssetCategory.CHARACTERS,
|
|
description="Test character with custom grooming task"
|
|
)
|
|
db.add(asset)
|
|
db.commit()
|
|
db.refresh(asset)
|
|
print(f" ✓ Created asset: {asset.name} (ID: {asset.id})")
|
|
|
|
# Create standard and custom tasks for the asset
|
|
standard_tasks = ["modeling", "surfacing", "rigging"]
|
|
custom_tasks = ["grooming", "lookdev"]
|
|
|
|
for task_type in standard_tasks + custom_tasks:
|
|
task = Task(
|
|
project_id=project.id,
|
|
asset_id=asset.id,
|
|
task_type=task_type,
|
|
name=f"{asset.name} - {task_type.title()}",
|
|
description=f"{task_type.title()} task for {asset.name}"
|
|
)
|
|
db.add(task)
|
|
|
|
db.commit()
|
|
print(f" ✓ Created {len(standard_tasks)} standard tasks")
|
|
print(f" ✓ Created {len(custom_tasks)} custom tasks")
|
|
|
|
# 3. Verify tasks were created
|
|
print("\n3. Verifying asset tasks...")
|
|
asset_tasks = db.query(Task).filter(Task.asset_id == asset.id).all()
|
|
print(f" ✓ Total tasks: {len(asset_tasks)}")
|
|
for task in asset_tasks:
|
|
task_label = "CUSTOM" if task.task_type in custom_tasks else "STANDARD"
|
|
print(f" - {task.name} ({task_label})")
|
|
|
|
# 4. Create an episode and shot with custom task type
|
|
print("\n4. Creating episode and shot with custom task type...")
|
|
episode = Episode(
|
|
project_id=project.id,
|
|
name="Episode 01",
|
|
episode_number=1
|
|
)
|
|
db.add(episode)
|
|
db.commit()
|
|
db.refresh(episode)
|
|
print(f" ✓ Created episode: {episode.name} (ID: {episode.id})")
|
|
|
|
shot = Shot(
|
|
episode_id=episode.id,
|
|
name="SH010",
|
|
frame_start=1001,
|
|
frame_end=1100
|
|
)
|
|
db.add(shot)
|
|
db.commit()
|
|
db.refresh(shot)
|
|
print(f" ✓ Created shot: {shot.name} (ID: {shot.id})")
|
|
|
|
# Create standard and custom tasks for the shot
|
|
standard_shot_tasks = ["layout", "animation", "lighting", "compositing"]
|
|
custom_shot_tasks = ["previz", "matchmove"]
|
|
|
|
for task_type in standard_shot_tasks + custom_shot_tasks:
|
|
task = Task(
|
|
project_id=project.id,
|
|
episode_id=episode.id,
|
|
shot_id=shot.id,
|
|
task_type=task_type,
|
|
name=f"{shot.name}_{task_type}",
|
|
description=f"{task_type.title()} task for shot {shot.name}"
|
|
)
|
|
db.add(task)
|
|
|
|
db.commit()
|
|
print(f" ✓ Created {len(standard_shot_tasks)} standard tasks")
|
|
print(f" ✓ Created {len(custom_shot_tasks)} custom tasks")
|
|
|
|
# 5. Verify shot tasks were created
|
|
print("\n5. Verifying shot tasks...")
|
|
shot_tasks = db.query(Task).filter(Task.shot_id == shot.id).all()
|
|
print(f" ✓ Total tasks: {len(shot_tasks)}")
|
|
for task in shot_tasks:
|
|
task_label = "CUSTOM" if task.task_type in custom_shot_tasks else "STANDARD"
|
|
print(f" - {task.name} ({task_label})")
|
|
|
|
# 6. Test querying all task types for the project
|
|
print("\n6. Testing task type queries...")
|
|
from routers.assets import get_all_asset_task_types
|
|
from routers.shots import get_all_shot_task_types
|
|
|
|
all_asset_types = get_all_asset_task_types(project.id, db)
|
|
all_shot_types = get_all_shot_task_types(project.id, db)
|
|
|
|
print(f" ✓ All asset task types: {all_asset_types}")
|
|
print(f" ✓ All shot task types: {all_shot_types}")
|
|
|
|
# Verify custom types are included
|
|
assert "grooming" in all_asset_types, "Custom asset type 'grooming' not found"
|
|
assert "lookdev" in all_asset_types, "Custom asset type 'lookdev' not found"
|
|
assert "previz" in all_shot_types, "Custom shot type 'previz' not found"
|
|
assert "matchmove" in all_shot_types, "Custom shot type 'matchmove' not found"
|
|
print(" ✓ All custom types are included in queries")
|
|
|
|
print("\n=== All Tests Passed! ===\n")
|
|
|
|
# Cleanup
|
|
print("Cleaning up test data...")
|
|
db.query(Task).filter(Task.project_id == project.id).delete()
|
|
db.query(Shot).filter(Shot.episode_id == episode.id).delete()
|
|
db.query(Episode).filter(Episode.project_id == project.id).delete()
|
|
db.query(Asset).filter(Asset.project_id == project.id).delete()
|
|
db.query(Project).filter(Project.id == project.id).delete()
|
|
db.commit()
|
|
print("✓ Cleanup complete")
|
|
|
|
except Exception as e:
|
|
print(f"\n❌ Error: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
db.rollback()
|
|
finally:
|
|
db.close()
|
|
|
|
if __name__ == "__main__":
|
|
test_custom_task_types_integration()
|