tests/test_generation_foundation.py

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,
        )