tests/test_generation_flux.py

"""Mocked protocol tests for the audited FLUX.2 Klein worker integration."""

from __future__ import annotations

from collections.abc import Callable
from pathlib import Path
from types import SimpleNamespace

import httpx
import pytest

from app.core.config import Settings
from app.generation.domain.enums import (
    GenerationModality,
    WorkerCancellationStatus,
    WorkerErrorCategory,
    WorkerJobStatus,
)
from app.generation.domain.errors import (
    GenerationCapabilityUnsupportedError,
    GenerationValidationError,
    GenerationWorkerError,
)
from app.generation.domain.retry import GenerationRetryPolicy
from app.generation.model_registry import GenerationModelRegistration, GenerationModelRegistry
from app.generation.providers.flux import (
    FLUX_BASE_MODEL_ID,
    FLUX_DISTILLED_MODEL_ID,
    FLUX_MODEL_CAPABILITY,
    FLUX_MODEL_ID,
    FLUX_PROVIDER_ID,
    FluxProviderAdapter,
)
from app.generation.providers.worker_client import RemoteWorkerClient
from app.generation.schemas.requests import GenerationRequestCreate


async def _no_sleep() -> None:
    return None


def _client(
    handler: Callable[[httpx.Request], httpx.Response], *, retries: int = 2
) -> RemoteWorkerClient:
    return RemoteWorkerClient(
        base_url="https://flux-worker.example",
        bearer_token="x" * 32,
        connect_timeout_seconds=1,
        request_timeout_seconds=1,
        read_timeout_seconds=1,
        retry_policy=GenerationRetryPolicy(max_retries=retries, backoff_seconds=0),
        http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)),
        sleep=lambda _: _no_sleep(),
    )


def _info() -> dict[str, object]:
    return {
        "id": FLUX_MODEL_ID,
        "name": "FLUX.2 Klein 4B",
        "type": "image",
        "license": "Apache-2.0",
        "status": "ready",
        "models": {"distilled": FLUX_DISTILLED_MODEL_ID, "base": FLUX_BASE_MODEL_ID},
    }


def _payload(**overrides: object) -> GenerationRequestCreate:
    value: dict[str, object] = {
        "provider": FLUX_PROVIDER_ID,
        "model_id": FLUX_MODEL_ID,
        "modality": "image",
        "prompt": "A cinematic coastal city at sunrise",
        "flux": {
            "mode_choice": "Distilled (4 steps)",
            "seed": 42,
            "randomize_seed": False,
            "width": 1024,
            "height": 1024,
            "num_inference_steps": 4,
            "guidance_scale": 1.0,
            "prompt_upsampling": False,
        },
    }
    value.update(overrides)
    return GenerationRequestCreate.model_validate(value)


@pytest.mark.asyncio
async def test_flux_exact_model_discovery_and_readiness() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        if request.url.path == "/health":
            return httpx.Response(200, json={"status": "ok"})
        if request.url.path == "/ready":
            return httpx.Response(
                200,
                json={
                    "status": "ready",
                    "model_loaded": True,
                    "model": FLUX_MODEL_ID,
                    "accepting_jobs": True,
                },
            )
        return httpx.Response(200, json=_info())

    adapter = FluxProviderAdapter(client=_client(handler))
    registry = GenerationModelRegistry(
        [
            GenerationModelRegistration(
                provider_id=FLUX_PROVIDER_ID,
                model=FLUX_MODEL_CAPABILITY,
                configuration_reference="flux-space",
            )
        ]
    )
    assert (await adapter.health()).status.value == "healthy"
    models = registry.verify_readiness(
        provider_id=FLUX_PROVIDER_ID,
        worker_info=await adapter.info(),
        readiness=await adapter.ready(),
        provider_configured=adapter.available,
    )
    assert models[0].model.id == FLUX_MODEL_ID
    assert models[0].model.modality is GenerationModality.IMAGE
    assert models[0].available


@pytest.mark.asyncio
async def test_flux_identity_mismatch_and_not_ready_are_not_advertised() -> None:
    wrong = {**_info(), "models": {"distilled": "untrusted/model", "base": FLUX_BASE_MODEL_ID}}

    def identity_handler(_: httpx.Request) -> httpx.Response:
        return httpx.Response(200, json=wrong)

    with pytest.raises(GenerationWorkerError) as raised:
        await FluxProviderAdapter(client=_client(identity_handler)).info()
    assert raised.value.category is WorkerErrorCategory.PROVIDER_ERROR

    def not_ready_handler(request: httpx.Request) -> httpx.Response:
        if request.url.path == "/ready":
            return httpx.Response(
                503,
                json={
                    "status": "not_ready",
                    "model_loaded": False,
                    "model": FLUX_MODEL_ID,
                    "accepting_jobs": False,
                },
            )
        return httpx.Response(200, json=_info())

    with pytest.raises(GenerationWorkerError) as raised:
        await FluxProviderAdapter(client=_client(not_ready_handler)).ready()
    assert raised.value.category is WorkerErrorCategory.WORKER_NOT_READY


@pytest.mark.asyncio
async def test_flux_text_submission_uses_strict_form_and_has_no_automatic_retry() -> None:
    requests: list[httpx.Request] = []

    def handler(request: httpx.Request) -> httpx.Response:
        requests.append(request)
        return httpx.Response(202, json={"job_id": "flux_" + "a" * 32, "status": "queued"})

    job = await FluxProviderAdapter(client=_client(handler)).submit(
        payload={"prompt": "A city at sunrise", "flux": {"width": 1024, "height": 1024}},
        idempotency_key="generation-request-id",
    )
    assert job.status is WorkerJobStatus.QUEUED
    assert requests[0].headers["authorization"] == "Bearer " + "x" * 32
    assert requests[0].headers["content-type"].startswith("application/x-www-form-urlencoded")
    assert b"width=1024" in requests[0].content


@pytest.mark.asyncio
async def test_flux_optional_canonical_image_uses_multipart(tmp_path: Path) -> None:
    source = tmp_path / "input.png"
    source.write_bytes(b"image-input")

    def handler(request: httpx.Request) -> httpx.Response:
        body = request.content.decode("latin-1")
        assert 'name="input_images"' in body
        assert 'name="prompt"' in body
        return httpx.Response(202, json={"job_id": "flux_" + "b" * 32, "status": "queued"})

    job = await FluxProviderAdapter(client=_client(handler)).submit(
        payload={"prompt": "Edit this image"},
        idempotency_key="generation-request-id",
        input_path=source,
        input_mime_type="image/png",
    )
    assert job.external_job_id.startswith("flux_")


@pytest.mark.asyncio
async def test_flux_rejects_invalid_requests_and_input_assets() -> None:
    adapter = FluxProviderAdapter(client=None)
    for invalid in ({"prompt": " "}, {"modality": "video"}):
        with pytest.raises(Exception):
            await adapter.validate_request(_payload(**invalid))

    with pytest.raises(GenerationValidationError):
        await adapter.validate_input_asset(
            _payload(), SimpleNamespace(mime_type="video/mp4", file_size=100)
        )
    with pytest.raises(GenerationValidationError):
        await adapter.validate_input_asset(
            _payload(), SimpleNamespace(mime_type="image/png", file_size=21 * 1024 * 1024)
        )


@pytest.mark.parametrize(
    "field,value",
    [
        ("negative_prompt", "unsupported"),
        ("scheduler", "unsupported"),
        ("width", 1023),
        ("height", 1032),
    ],
)
def test_flux_schema_rejects_unsupported_or_invalid_parameters(field: str, value: object) -> None:
    raw = _payload().model_dump()
    flux = dict(raw["flux"] or {})
    flux[field] = value
    raw["flux"] = flux
    with pytest.raises(ValueError):
        GenerationRequestCreate.model_validate(raw)


@pytest.mark.asyncio
async def test_flux_rejects_controls_for_another_provider() -> None:
    payload = _payload(wan={"duration_seconds": 1.0})
    with pytest.raises(GenerationCapabilityUnsupportedError):
        await FluxProviderAdapter(client=None).validate_request(payload)


@pytest.mark.asyncio
async def test_flux_completed_job_maps_a_safe_png_output_and_retrieves_it() -> None:
    job_id = "flux_" + "c" * 32

    def handler(request: httpx.Request) -> httpx.Response:
        if request.url.path.endswith("/output"):
            return httpx.Response(200, content=b"png-output")
        return httpx.Response(
            200,
            json={
                "job_id": job_id,
                "status": "completed",
                "output": {"type": "image", "filename": "output.png"},
            },
        )

    adapter = FluxProviderAdapter(client=_client(handler))
    job = await adapter.get_job(external_job_id=job_id)
    assert job.output is not None
    assert job.output.mime_type == "image/png"
    assert job.output.download_path == f"/v1/jobs/{job_id}/output"
    output = await adapter.retrieve_output(external_job_id=job_id)
    async with adapter.stream_output(output) as chunks:
        assert b"".join([chunk async for chunk in chunks]) == b"png-output"


@pytest.mark.asyncio
@pytest.mark.parametrize("status_code", [429, 502, 503, 504])
async def test_flux_polling_uses_shared_bounded_transient_retry(status_code: int) -> None:
    calls = 0
    job_id = "flux_" + "d" * 32

    def handler(_: httpx.Request) -> httpx.Response:
        nonlocal calls
        calls += 1
        if calls < 3:
            return httpx.Response(status_code, json={"detail": {"token": "never-store"}})
        return httpx.Response(200, json={"job_id": job_id, "status": "running"})

    job = await FluxProviderAdapter(client=_client(handler, retries=2)).get_job(
        external_job_id=job_id
    )
    assert job.status is WorkerJobStatus.RUNNING
    assert calls == 3


@pytest.mark.asyncio
async def test_flux_permanent_error_is_not_retried_and_cancellation_is_accurate() -> None:
    job_id = "flux_" + "e" * 32
    calls = 0

    def permanent_handler(_: httpx.Request) -> httpx.Response:
        nonlocal calls
        calls += 1
        return httpx.Response(400, json={"detail": {"code": "FLUX_REQUEST_INVALID"}})

    with pytest.raises(GenerationWorkerError) as raised:
        await FluxProviderAdapter(client=_client(permanent_handler, retries=3)).get_job(
            external_job_id=job_id
        )
    assert raised.value.category is WorkerErrorCategory.INVALID_REQUEST
    assert calls == 1

    def queued_handler(_: httpx.Request) -> httpx.Response:
        return httpx.Response(200, json={"job_id": job_id, "status": "cancelled"})

    def running_handler(_: httpx.Request) -> httpx.Response:
        return httpx.Response(
            409,
            json={"detail": {"code": "FLUX_JOB_NOT_CANCELLABLE", "status": "running"}},
        )

    assert (
        await FluxProviderAdapter(client=_client(queued_handler)).cancel(external_job_id=job_id)
    ).status is WorkerCancellationStatus.CANCELLED
    assert (
        await FluxProviderAdapter(client=_client(running_handler)).cancel(external_job_id=job_id)
    ).status is WorkerCancellationStatus.FAILED


def test_flux_configuration_is_optional_and_does_not_change_wan_configuration() -> None:
    disabled = FluxProviderAdapter.from_settings(Settings(_env_file=None))
    invalid = FluxProviderAdapter.from_settings(
        Settings(_env_file=None, flux_space_url="https://flux-worker.example")
    )
    assert not disabled.available
    assert not invalid.available
    assert invalid.configuration_error is not None