feat: add resource library and real image workflow
This commit is contained in:
afei A
298
tests/test_tryon_activity.py
Normal file
298
tests/test_tryon_activity.py
Normal file
@@ -0,0 +1,298 @@
|
||||
"""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"
|
||||
Reference in New Issue
Block a user