"""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