from __future__ import annotations import base64 from contextlib import asynccontextmanager from pathlib import Path import pytest from fastapi.testclient import TestClient from app.container import build_container from app.ai.schemas import AiGenerateImageRequest from app.core.config import Settings from app.generation.domain.capabilities import ( GenerationModelCapability, GenerationProviderCapabilities, ) from app.generation.domain.enums import ( GenerationJobStatus, GenerationModality, WorkerCancellationStatus, WorkerHealthStatus, WorkerJobStatus, WorkerReadinessStatus, ) from app.generation.domain.errors import ( GenerationIdempotencyConflictError, GenerationInputAssetNotFoundError, GenerationJobNotFoundError, GenerationProviderJobConflictError, ) from app.generation.domain.runtime import ( WorkerCancellationResult, WorkerHealth, WorkerInfo, WorkerJob, WorkerOutput, WorkerReadiness, ) from app.generation.model_registry import ( GenerationModelRegistration, GenerationModelRegistry, ) from app.generation.providers.base import GenerationProviderAdapter from app.generation.providers.registry import GenerationProviderRegistry from app.generation.schemas.requests import GenerationRequestCreate from app.security.schemas import APIKeyCreate from main import create_app def generation_settings(tmp_path: Path) -> Settings: return Settings( _env_file=None, auth_enabled=True, database_url=f"sqlite+aiosqlite:///{tmp_path / 'security.db'}", social_database_url=f"sqlite+aiosqlite:///{tmp_path / 'social.db'}", social_auto_migrate=True, social_worker_enabled=False, social_oauth_encryption_key="test-only-encryption-material", temp_dir=tmp_path / "temp", output_dir=tmp_path / "outputs", cleanup_interval_seconds=3600, whisper_model="tiny", generation_enabled=True, ) class AvailableTestProvider(GenerationProviderAdapter): capabilities = GenerationProviderCapabilities( provider="test-generation", name="Test generation adapter", implementation_status="test", models=[ GenerationModelCapability( id="test-image-v1", name="Test image v1", modality=GenerationModality.IMAGE, input_asset_supported=True, ) ], ) def __init__(self) -> None: self.cancellation_result = WorkerCancellationResult( status=WorkerCancellationStatus.REQUESTED ) @property def available(self) -> bool: return True async def validate_request(self, payload: GenerationRequestCreate) -> dict[str, object]: return {"prompt": payload.prompt} async def health(self) -> WorkerHealth: return WorkerHealth(status=WorkerHealthStatus.HEALTHY) async def info(self) -> WorkerInfo: return WorkerInfo( id="test-generation-worker", name="Test generation worker", media_types=[GenerationModality.IMAGE], models=[ { "id": "test-image-v1", "name": "Test image v1", "media_types": [GenerationModality.IMAGE], } ], status=WorkerHealthStatus.HEALTHY, ) async def ready(self) -> WorkerReadiness: return WorkerReadiness( status=WorkerReadinessStatus.READY, model_loaded=True, model_ids=["test-image-v1"], ) async def cancel(self, *, external_job_id: str) -> WorkerCancellationResult: assert external_job_id == "worker-job-1" return self.cancellation_result async def get_job(self, *, external_job_id: str) -> WorkerJob: assert external_job_id == "worker-job-1" return WorkerJob( external_job_id=external_job_id, status=WorkerJobStatus.COMPLETED, output=WorkerOutput( output_type=GenerationModality.IMAGE, mime_type="image/png", provider_output_id="worker-output-1", download_path="/v1/outputs/worker-output-1", ), ) @asynccontextmanager async def stream_output(self, output: WorkerOutput): assert output.provider_output_id == "worker-output-1" async def chunks(): yield base64.b64decode( "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ" "AAAADUlEQVQIHWP4z8DwHwAFgAI/ScL9aQAAAABJRU5ErkJggg==" ) yield chunks() async def create_context(container, name: str): _, secret = await container.api_keys.create( APIKeyCreate( name=name, environment="test", role=None, scopes=[ "generation:providers:read", "generation:requests:read", "generation:requests:create", "generation:jobs:cancel", ], ), created_by="tests", ) return await container.api_keys.authenticate(secret) def request_payload(*, prompt: str = "A test image") -> GenerationRequestCreate: return GenerationRequestCreate( provider="test-generation", model_id="test-image-v1", modality=GenerationModality.IMAGE, prompt=prompt, ) @pytest.fixture async def generation_container(tmp_path: Path): container = build_container(generation_settings(tmp_path)) await container.security_database.initialize() provider = AvailableTestProvider() container.generation.providers = GenerationProviderRegistry([provider]) container.generation.models = GenerationModelRegistry( [ GenerationModelRegistration( provider_id=provider.provider, model=provider.capabilities.models[0], configuration_reference="test-generation-worker", ) ] ) await container.generation.initialize() await container.generation.refresh_provider_runtime(provider.provider) try: yield container finally: await container.security_database.close() @pytest.mark.asyncio async def test_optional_generation_providers_start_unavailable_without_configuration( tmp_path: Path, ) -> None: container = build_container(generation_settings(tmp_path)) await container.security_database.initialize() await container.generation.initialize() try: providers = container.generation.list_providers() assert [provider.capabilities.provider for provider in providers] == ["flux", "wan"] assert not any(provider.available for provider in providers) assert not container.generation.get_model("flux", "flux.2-klein-4b").available assert not container.generation.get_model("wan", "wan2.2").available finally: await container.security_database.close() @pytest.mark.asyncio async def test_ai_studio_advertises_and_isolates_real_generation_history( generation_container, ) -> None: context = await create_context(generation_container, "AI Studio") capabilities = generation_container.ai.capabilities() image_tool = next(tool for tool in capabilities.tools if tool.operation == "generate_image") video_tool = next(tool for tool in capabilities.tools if tool.operation == "generate_video") assert image_tool.available assert not video_tool.available ordinary = await generation_container.generation.create( workspace_id=context.workspace_id, user_id=context.user_id, payload=request_payload(prompt="ordinary generation"), idempotency_key="ordinary-generation-key", ) ai_job = await generation_container.ai.create( workspace_id=context.workspace_id, user_id=context.user_id, api_key_id=context.api_key_id, request_id="ai-request", payload=AiGenerateImageRequest( operation="generate_image", prompt="AI Studio generation", ), idempotency_key="ai-studio-generation-key", ) history = await generation_container.ai.history( workspace_id=context.workspace_id, user_id=context.user_id, offset=0, limit=25, ) assert [item.generation_id for item in history.items] == [ai_job.generation_id] assert ordinary.id not in {item.generation_id for item in history.items} def test_application_starts_with_optional_providers_disabled_when_unconfigured( tmp_path: Path, ) -> None: """No worker URL/token is needed merely to start the application.""" with TestClient(create_app(generation_settings(tmp_path))) as client: providers = client.app.state.container.generation.list_providers() assert [provider.capabilities.provider for provider in providers] == ["flux", "wan"] models = client.app.state.container.generation.list_models() assert [model.model.id for model in models] == ["flux.2-klein-4b", "wan2.2"] assert not any(model.available for model in models) @pytest.mark.asyncio async def test_provider_discovery_requires_a_verified_model(generation_container) -> None: """A configured adapter is not publicly usable before runtime verification.""" provider_id = "test-generation" generation_container.generation.models.mark_unavailable(provider_id) assert not generation_container.generation.get_provider(provider_id).available assert not generation_container.generation.list_providers()[0].available await generation_container.generation.refresh_provider_runtime(provider_id) assert generation_container.generation.get_provider(provider_id).available @pytest.mark.asyncio async def test_generation_request_idempotency_and_cancel(generation_container) -> None: context = await create_context(generation_container, "Generation A") workspace_id = str(context.workspace_id) user_id = str(context.user_id) first = await generation_container.generation.create( workspace_id=workspace_id, user_id=user_id, payload=request_payload(), idempotency_key="generation-request-key", ) replay = await generation_container.generation.create( workspace_id=workspace_id, user_id=user_id, payload=request_payload(), idempotency_key="generation-request-key", ) assert replay.id == first.id assert replay.job.id == first.job.id with pytest.raises(GenerationIdempotencyConflictError): await generation_container.generation.create( workspace_id=workspace_id, user_id=user_id, payload=request_payload(prompt="Different request"), idempotency_key="generation-request-key", ) cancelled = await generation_container.generation.cancel(workspace_id, user_id, first.job.id) assert cancelled.status is GenerationJobStatus.CANCELLED retrieved = await generation_container.generation.get_request(workspace_id, user_id, first.id) assert retrieved.status is GenerationJobStatus.CANCELLED @pytest.mark.asyncio async def test_generation_records_are_workspace_isolated(generation_container) -> None: context_a = await create_context(generation_container, "Generation A") context_b = await create_context(generation_container, "Generation B") created = await generation_container.generation.create( workspace_id=str(context_a.workspace_id), user_id=str(context_a.user_id), payload=request_payload(), idempotency_key="generation-isolation-key", ) with pytest.raises(GenerationJobNotFoundError): await generation_container.generation.get_job( str(context_b.workspace_id), str(context_b.user_id), created.job.id ) assert ( await generation_container.generation.list_requests( str(context_b.workspace_id), str(context_b.user_id) ) == [] ) @pytest.mark.asyncio async def test_generation_rejects_another_workspace_canonical_input_asset( generation_container, ) -> None: context_a = await create_context(generation_container, "Generation A") context_b = await create_context(generation_container, "Generation B") request_id = "00000000-0000-0000-0000-000000000010" output_dir = generation_container.settings.output_dir / request_id output_dir.mkdir(parents=True) output = output_dir / "owned-input.png" output.write_bytes(b"canonical image") asset = await generation_container.assets.register_output( workspace_id=str(context_a.workspace_id), user_id=str(context_a.user_id), request_id=request_id, path=output, mime_type="image/png", ) with pytest.raises(GenerationInputAssetNotFoundError): await generation_container.generation.create( workspace_id=str(context_b.workspace_id), user_id=str(context_b.user_id), payload=GenerationRequestCreate( provider="test-generation", model_id="test-image-v1", modality=GenerationModality.IMAGE, prompt="Use another workspace asset", input_asset_id=asset.id, ), idempotency_key="generation-cross-asset-key", ) @pytest.mark.parametrize("forbidden_field", ["provider_payload", "worker_url", "output_url"]) def test_generation_request_schema_rejects_client_supplied_provider_controls( forbidden_field: str, ) -> None: payload: dict[str, object] = { "provider": "test-generation", "model_id": "test-image-v1", "modality": "image", "prompt": "A test image", } payload[forbidden_field] = {"unsafe": True} with pytest.raises(ValueError): GenerationRequestCreate.model_validate(payload) @pytest.mark.asyncio async def test_remote_cancellation_preserves_requested_and_confirmed_states( generation_container, ) -> None: context = await create_context(generation_container, "Generation cancellation") workspace_id = str(context.workspace_id) user_id = str(context.user_id) created = await generation_container.generation.create( workspace_id=workspace_id, user_id=user_id, payload=request_payload(), idempotency_key="generation-cancellation-key", ) await generation_container.generation.repository.transition_job( workspace_id, created.job.id, GenerationJobStatus.SUBMITTING, user_id=user_id, ) await generation_container.generation.bind_provider_job( workspace_id=workspace_id, user_id=user_id, job_id=created.job.id, worker_job_id="worker-job-1", ) await generation_container.generation.repository.transition_job( workspace_id, created.job.id, GenerationJobStatus.RUNNING, user_id=user_id, ) requested = await generation_container.generation.cancel(workspace_id, user_id, created.job.id) assert requested.status is GenerationJobStatus.CANCEL_REQUESTED provider = generation_container.generation.providers.get("test-generation") assert isinstance(provider, AvailableTestProvider) provider.cancellation_result = WorkerCancellationResult( status=WorkerCancellationStatus.CANCELLED ) confirmed = await generation_container.generation.cancel(workspace_id, user_id, created.job.id) assert confirmed.status is GenerationJobStatus.CANCELLED @pytest.mark.asyncio async def test_provider_job_binding_and_output_ingestion_are_workspace_scoped( generation_container, ) -> None: context_a = await create_context(generation_container, "Generation output A") context_b = await create_context(generation_container, "Generation output B") workspace_a, user_a = str(context_a.workspace_id), str(context_a.user_id) workspace_b, user_b = str(context_b.workspace_id), str(context_b.user_id) job_a = await generation_container.generation.create( workspace_id=workspace_a, user_id=user_a, payload=request_payload(), idempotency_key="generation-output-a", ) job_b = await generation_container.generation.create( workspace_id=workspace_b, user_id=user_b, payload=request_payload(), idempotency_key="generation-output-b", ) for workspace_id, user_id, job_id in ( (workspace_a, user_a, job_a.job.id), (workspace_b, user_b, job_b.job.id), ): await generation_container.generation.repository.transition_job( workspace_id, job_id, GenerationJobStatus.SUBMITTING, user_id=user_id, ) await generation_container.generation.bind_provider_job( workspace_id=workspace_a, user_id=user_a, job_id=job_a.job.id, worker_job_id="worker-job-1", ) with pytest.raises(GenerationProviderJobConflictError): await generation_container.generation.bind_provider_job( workspace_id=workspace_b, user_id=user_b, job_id=job_b.job.id, worker_job_id="worker-job-1", ) await generation_container.generation.repository.transition_job( workspace_a, job_a.job.id, GenerationJobStatus.RUNNING, user_id=user_a, ) completed = await generation_container.generation.ingest_completed_provider_output( workspace_id=workspace_a, user_id=user_a, job_id=job_a.job.id, ) assert completed.status is GenerationJobStatus.SUCCEEDED assert completed.output_asset_id is not None output_asset = await generation_container.assets.get_owned_by_id( workspace_id=workspace_a, user_id=user_a, asset_id=completed.output_asset_id, ) assert output_asset.mime_type == "image/png" assert output_asset.metadata_json["generation"]["media"]["resolution"] == { "width": 1, "height": 1, } assert ( await generation_container.generation.ingest_completed_provider_output( workspace_id=workspace_a, user_id=user_a, job_id=job_a.job.id, ) == completed ) with pytest.raises(GenerationJobNotFoundError): await generation_container.generation.ingest_completed_provider_output( workspace_id=workspace_b, user_id=user_b, job_id=job_a.job.id, )