Spaces:
Running
Running
| 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 ( | |
| GenerationValidationError, | |
| 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 app.projects.models import Project | |
| 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 | |
| ) | |
| 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", | |
| ), | |
| ) | |
| 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, | |
| ) | |
| 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() | |
| 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() | |
| 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) | |
| 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 | |
| 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 | |
| async def test_failed_generation_job_can_be_retried_as_new_durable_request( | |
| generation_container, | |
| ) -> None: | |
| context = await create_context(generation_container, "Generation retry") | |
| original = await generation_container.generation.create( | |
| workspace_id=context.workspace_id, | |
| user_id=context.user_id, | |
| payload=request_payload(prompt="Retry this image"), | |
| idempotency_key="generation-retry-original", | |
| ) | |
| await generation_container.generation.transition_worker_job( | |
| workspace_id=context.workspace_id, | |
| user_id=context.user_id, | |
| job_id=original.job.id, | |
| status=GenerationJobStatus.FAILED, | |
| error_code="TEST_FAILURE", | |
| error_message="Expected test failure", | |
| ) | |
| retried = await generation_container.generation.retry_failed_job( | |
| context.workspace_id, | |
| context.user_id, | |
| original.job.id, | |
| idempotency_key="generation-retry-new-request", | |
| ) | |
| assert retried.id != original.id | |
| assert retried.job.id != original.job.id | |
| assert retried.prompt == original.prompt | |
| assert retried.provider == original.provider | |
| assert retried.model_id == original.model_id | |
| assert retried.status.value == "queued" | |
| async def test_generation_attempt_history_is_durable_and_workspace_scoped( | |
| generation_container, | |
| ) -> None: | |
| context = await create_context(generation_container, "Generation attempt history") | |
| request = await generation_container.generation.create( | |
| workspace_id=context.workspace_id, | |
| user_id=context.user_id, | |
| payload=request_payload(), | |
| idempotency_key="generation-attempt-history", | |
| ) | |
| claimed = await generation_container.generation.claim_dispatchable_jobs(limit=1) | |
| assert len(claimed) == 1 | |
| assert claimed[0].generation_job_id == request.job.id | |
| attempts = await generation_container.generation.list_job_attempts( | |
| context.workspace_id, context.user_id, request.job.id | |
| ) | |
| assert len(attempts) == 1 | |
| assert attempts[0].attempt_number == 1 | |
| assert attempts[0].status == "submitting" | |
| assert attempts[0].generation_job_id == request.job.id | |
| other_context = await create_context(generation_container, "Other workspace") | |
| with pytest.raises(GenerationJobNotFoundError): | |
| await generation_container.generation.list_job_attempts( | |
| other_context.workspace_id, other_context.user_id, request.job.id | |
| ) | |
| 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) | |
| ) | |
| == [] | |
| ) | |
| 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", | |
| ) | |
| 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) | |
| 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 | |
| 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, | |
| ) | |
| async def test_generation_request_attaches_only_to_active_project_in_same_workspace(generation_container) -> None: | |
| container = generation_container | |
| context_a = await create_context(container, "Project generation workspace A") | |
| context_b = await create_context(container, "Project generation workspace B") | |
| async def add_project(context, name: str) -> Project: | |
| async with container.security_database.tenant_session(workspace_id=context.workspace_id, user_id=context.user_id) as session: | |
| project = Project(workspace_id=context.workspace_id, created_by=context.user_id, name=name, status="active") | |
| session.add(project) | |
| await session.commit() | |
| await session.refresh(project) | |
| return project | |
| owned_project = await add_project(context_a, "Owned generation project") | |
| other_project = await add_project(context_b, "Other workspace project") | |
| payload = request_payload().model_copy(update={"project_id": owned_project.id}) | |
| created = await container.generation.create(workspace_id=context_a.workspace_id, user_id=context_a.user_id, payload=payload, idempotency_key="project-link-owned-001") | |
| assert created.project_id == owned_project.id | |
| foreign_payload = request_payload().model_copy(update={"project_id": other_project.id}) | |
| with pytest.raises(GenerationValidationError, match="Project was not found"): | |
| await container.generation.create(workspace_id=context_a.workspace_id, user_id=context_a.user_id, payload=foreign_payload, idempotency_key="project-link-foreign-001") | |
| def test_generation_request_schema_accepts_only_uuid_project_ids() -> None: | |
| from pydantic import ValidationError | |
| base = {"provider": "test-generation", "model_id": "test-image-v1", "modality": "image", "prompt": "A test image"} | |
| valid = GenerationRequestCreate.model_validate({**base, "project_id": "123e4567-e89b-42d3-a456-426614174000"}) | |
| assert valid.project_id == "123e4567-e89b-42d3-a456-426614174000" | |
| with pytest.raises(ValidationError): | |
| GenerationRequestCreate.model_validate({**base, "project_id": "../other-workspace"}) | |