"""Focused tests for the try-on activity implementation.""" from __future__ import annotations from dataclasses import dataclass import pytest from app.domain.enums import AssetType, LibraryFileRole, LibraryResourceStatus, LibraryResourceType, OrderStatus, WorkflowStepName from app.infra.db.models.asset import AssetORM from app.infra.db.models.library_resource import LibraryResourceORM from app.infra.db.models.library_resource_file import LibraryResourceFileORM from app.infra.db.models.order import OrderORM from app.infra.db.models.workflow_run import WorkflowRunORM from app.infra.db.session import get_session_factory from app.workers.activities import tryon_activities from app.workers.workflows.types import StepActivityInput @dataclass(slots=True) class FakeGeneratedImage: image_bytes: bytes mime_type: str provider: str model: str prompt: str class FakeImageGenerationService: async def generate_tryon_image(self, *, person_image_url: str, garment_image_url: str) -> FakeGeneratedImage: assert person_image_url == "https://images.example.com/orders/1/prepared-model.png" assert garment_image_url == "https://images.example.com/library/garments/cream-dress/original.png" return FakeGeneratedImage( image_bytes=b"fake-png-binary", mime_type="image/png", provider="gemini", model="gemini-test-image", prompt="test prompt", ) async def generate_scene_image(self, *, source_image_url: str, scene_image_url: str) -> FakeGeneratedImage: assert source_image_url == "https://images.example.com/orders/1/tryon/generated.png" assert scene_image_url == "https://images.example.com/library/scenes/studio/original.png" return FakeGeneratedImage( image_bytes=b"fake-scene-binary", mime_type="image/jpeg", provider="gemini", model="gemini-test-image", prompt="scene prompt", ) class FakeOrderArtifactStorageService: async def upload_generated_image( self, *, order_id: int, step_name: WorkflowStepName, image_bytes: bytes, mime_type: str, ) -> tuple[str, str]: assert order_id == 1 assert step_name == WorkflowStepName.TRYON assert image_bytes == b"fake-png-binary" assert mime_type == "image/png" return ( "orders/1/tryon/generated.png", "https://images.example.com/orders/1/tryon/generated.png", ) class FakeSceneArtifactStorageService: async def upload_generated_image( self, *, order_id: int, step_name: WorkflowStepName, image_bytes: bytes, mime_type: str, ) -> tuple[str, str]: assert order_id == 1 assert step_name == WorkflowStepName.SCENE assert image_bytes == b"fake-scene-binary" assert mime_type == "image/jpeg" return ( "orders/1/scene/generated.jpg", "https://images.example.com/orders/1/scene/generated.jpg", ) @pytest.mark.asyncio async def test_run_tryon_activity_persists_uploaded_asset_with_provider_metadata(api_runtime, monkeypatch): """Gemini-mode try-on should persist the uploaded output URL instead of a mock URI.""" monkeypatch.setattr( tryon_activities, "get_image_generation_service", lambda: FakeImageGenerationService(), ) monkeypatch.setattr( tryon_activities, "get_order_artifact_storage_service", lambda: FakeOrderArtifactStorageService(), ) async with get_session_factory()() as session: order = OrderORM( customer_level="low", service_mode="auto_basic", status=OrderStatus.CREATED, model_id=1, garment_asset_id=2, ) session.add(order) await session.flush() workflow_run = WorkflowRunORM( order_id=order.id, workflow_id=f"order-{order.id}", workflow_type="LowEndPipelineWorkflow", status=OrderStatus.CREATED, ) session.add(workflow_run) await session.flush() prepared_asset = AssetORM( order_id=order.id, asset_type=AssetType.PREPARED_MODEL, step_name=WorkflowStepName.PREPARE_MODEL, uri="https://images.example.com/orders/1/prepared-model.png", metadata_json={"library_resource_id": 1}, ) session.add(prepared_asset) garment_resource = LibraryResourceORM( resource_type=LibraryResourceType.GARMENT, name="Cream Dress", description="米白色连衣裙", tags=["女装"], status=LibraryResourceStatus.ACTIVE, category="dress", ) session.add(garment_resource) await session.flush() garment_original = LibraryResourceFileORM( resource_id=garment_resource.id, file_role=LibraryFileRole.ORIGINAL, storage_key="library/garments/cream-dress/original.png", public_url="https://images.example.com/library/garments/cream-dress/original.png", bucket="test-bucket", mime_type="image/png", size_bytes=1024, sort_order=0, ) garment_thumb = LibraryResourceFileORM( resource_id=garment_resource.id, file_role=LibraryFileRole.THUMBNAIL, storage_key="library/garments/cream-dress/thumb.png", public_url="https://images.example.com/library/garments/cream-dress/thumb.png", bucket="test-bucket", mime_type="image/png", size_bytes=256, sort_order=0, ) session.add_all([garment_original, garment_thumb]) await session.flush() garment_resource.original_file_id = garment_original.id garment_resource.cover_file_id = garment_thumb.id await session.commit() payload = StepActivityInput( order_id=1, workflow_run_id=1, step_name=WorkflowStepName.TRYON, source_asset_id=1, garment_asset_id=1, ) result = await tryon_activities.run_tryon_activity(payload) assert result.uri == "https://images.example.com/orders/1/tryon/generated.png" assert result.metadata["provider"] == "gemini" assert result.metadata["model"] == "gemini-test-image" assert result.metadata["prepared_asset_id"] == 1 assert result.metadata["garment_resource_id"] == 1 async with get_session_factory()() as session: assets = (await session.execute( AssetORM.__table__.select().where(AssetORM.order_id == 1, AssetORM.asset_type == AssetType.TRYON) )).mappings().all() assert len(assets) == 1 assert assets[0]["uri"] == "https://images.example.com/orders/1/tryon/generated.png" @pytest.mark.asyncio async def test_run_scene_activity_persists_uploaded_asset_with_provider_metadata(api_runtime, monkeypatch): """Gemini-mode scene should persist the uploaded output URL instead of a mock URI.""" from app.workers.activities import scene_activities monkeypatch.setattr( scene_activities, "get_image_generation_service", lambda: FakeImageGenerationService(), ) monkeypatch.setattr( scene_activities, "get_order_artifact_storage_service", lambda: FakeSceneArtifactStorageService(), ) async with get_session_factory()() as session: order = OrderORM( customer_level="low", service_mode="auto_basic", status=OrderStatus.CREATED, model_id=1, garment_asset_id=2, scene_ref_asset_id=3, ) session.add(order) await session.flush() workflow_run = WorkflowRunORM( order_id=order.id, workflow_id=f"order-{order.id}", workflow_type="LowEndPipelineWorkflow", status=OrderStatus.CREATED, ) session.add(workflow_run) await session.flush() tryon_asset = AssetORM( order_id=order.id, asset_type=AssetType.TRYON, step_name=WorkflowStepName.TRYON, uri="https://images.example.com/orders/1/tryon/generated.png", metadata_json={"prepared_asset_id": 1}, ) session.add(tryon_asset) scene_resource = LibraryResourceORM( resource_type=LibraryResourceType.SCENE, name="Studio Background", description="摄影棚背景", tags=["室内"], status=LibraryResourceStatus.ACTIVE, environment="indoor", ) session.add(scene_resource) await session.flush() scene_original = LibraryResourceFileORM( resource_id=scene_resource.id, file_role=LibraryFileRole.ORIGINAL, storage_key="library/scenes/studio/original.png", public_url="https://images.example.com/library/scenes/studio/original.png", bucket="test-bucket", mime_type="image/png", size_bytes=2048, sort_order=0, ) session.add(scene_original) await session.flush() scene_resource.original_file_id = scene_original.id scene_resource.cover_file_id = scene_original.id await session.commit() payload = StepActivityInput( order_id=1, workflow_run_id=1, step_name=WorkflowStepName.SCENE, source_asset_id=1, scene_ref_asset_id=1, ) result = await scene_activities.run_scene_activity(payload) assert result.uri == "https://images.example.com/orders/1/scene/generated.jpg" assert result.metadata["provider"] == "gemini" assert result.metadata["model"] == "gemini-test-image" assert result.metadata["source_asset_id"] == 1 assert result.metadata["scene_resource_id"] == 1 async with get_session_factory()() as session: assets = ( await session.execute( AssetORM.__table__.select().where( AssetORM.order_id == 1, AssetORM.asset_type == AssetType.SCENE, ) ) ).mappings().all() assert len(assets) == 1 assert assets[0]["uri"] == "https://images.example.com/orders/1/scene/generated.jpg"