tests/test_generation_provider_runtime.py

from __future__ import annotations

from collections.abc import Callable

import httpx
import pytest
from pydantic import ValidationError

from app.generation.domain.capabilities import (
    GenerationModelCapability,
    GenerationProviderCapabilities,
)
from app.generation.domain.enums import (
    GenerationModality,
    WorkerCancellationStatus,
    WorkerErrorCategory,
    WorkerHealthStatus,
    WorkerReadinessStatus,
)
from app.generation.domain.errors import GenerationWorkerError
from app.generation.domain.retry import GenerationRetryPolicy
from app.generation.domain.runtime import WorkerInfo, 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.providers.worker_client import RemoteWorkerClient


def worker_client(
    handler: Callable[[httpx.Request], httpx.Response] | None = None,
    *,
    retries: int = 2,
    sleep_calls: list[float] | None = None,
) -> RemoteWorkerClient:
    async def sleep(delay: float) -> None:
        if sleep_calls is not None:
            sleep_calls.append(delay)

    client = httpx.AsyncClient(
        transport=httpx.MockTransport(
            handler
            or (lambda _: httpx.Response(200, json={"status": "ok"}))
        )
    )
    return RemoteWorkerClient(
        base_url="https://worker.example",
        bearer_token="test-worker-token",
        connect_timeout_seconds=1,
        request_timeout_seconds=1,
        read_timeout_seconds=1,
        retry_policy=GenerationRetryPolicy(max_retries=retries, backoff_seconds=0),
        http_client=client,
        sleep=sleep,
    )


class RuntimeTestProvider(GenerationProviderAdapter):
    capabilities = GenerationProviderCapabilities(
        provider="runtime-test",
        name="Runtime test provider",
        implementation_status="test",
        models=[
            GenerationModelCapability(
                id="runtime-image-v1",
                name="Runtime image v1",
                modality=GenerationModality.IMAGE,
            )
        ],
    )


def test_provider_and_model_registration_starts_unavailable() -> None:
    provider = RuntimeTestProvider()
    providers = GenerationProviderRegistry([provider])
    assert providers.get("runtime-test") is provider
    models = GenerationModelRegistry(
        [
            GenerationModelRegistration(
                provider_id=provider.provider,
                model=provider.capabilities.models[0],
                configuration_reference="runtime-test-config",
                metadata={
                    "access_token": "must-not-survive",
                    "diagnostic": (
                        "Bearer must-not-survive "
                        "https://worker.example/output?sig=secret"
                    ),
                    "download_url": "https://worker.example/output?sig=secret",
                },
            )
        ]
    )
    view = models.get(provider.provider, "runtime-image-v1")
    assert not view.available
    assert "access_token" not in view.metadata
    assert "download_url" not in view.metadata
    assert "must-not-survive" not in str(view.metadata)


def test_model_availability_requires_readiness_info_and_configuration() -> None:
    model = GenerationModelCapability(
        id="runtime-image-v1", name="Runtime", modality=GenerationModality.IMAGE
    )
    registry = GenerationModelRegistry(
        [
            GenerationModelRegistration(
                provider_id="runtime-test",
                model=model,
                configuration_reference="runtime-test-config",
            )
        ]
    )
    info = WorkerInfo(
        id="runtime-test-worker",
        name="Runtime worker",
        media_types=[GenerationModality.IMAGE],
        models=[
            {
                "id": model.id,
                "name": model.name,
                "media_types": [GenerationModality.IMAGE],
            }
        ],
    )
    not_ready = WorkerReadiness(
        status=WorkerReadinessStatus.STARTING,
        model_loaded=False,
        model_ids=[model.id],
    )
    assert not registry.verify_readiness(
        provider_id="runtime-test",
        worker_info=info,
        readiness=not_ready,
        provider_configured=True,
    )[0].available
    ready = WorkerReadiness(
        status=WorkerReadinessStatus.READY, model_loaded=True, model_ids=[model.id]
    )
    assert registry.verify_readiness(
        provider_id="runtime-test",
        worker_info=info,
        readiness=ready,
        provider_configured=True,
    )[0].available


@pytest.mark.asyncio
async def test_worker_health_readiness_info_and_bearer_authentication() -> None:
    seen_headers: list[str] = []

    def handler(request: httpx.Request) -> httpx.Response:
        seen_headers.append(request.headers.get("authorization", ""))
        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": "model-v1"},
            )
        return httpx.Response(
            200,
            json={"id": "model-v1", "name": "Worker model", "type": "image", "status": "ready"},
        )

    client = worker_client(handler)
    assert (await client.health()).status is WorkerHealthStatus.HEALTHY
    readiness = await client.ready()
    assert readiness.status is WorkerReadinessStatus.READY
    assert readiness.model_ids == ["model-v1"]
    info = await client.info()
    assert info.media_types == [GenerationModality.IMAGE]
    assert info.models[0].id == "model-v1"
    assert seen_headers == ["Bearer test-worker-token"] * 3


@pytest.mark.asyncio
async def test_timeout_and_connection_failure_are_retryable_and_safe() -> None:
    request = httpx.Request("GET", "https://worker.example/health")
    for exception, category in (
        (httpx.ReadTimeout("secret-token", request=request), WorkerErrorCategory.TIMEOUT),
        (
            httpx.ConnectError("Bearer test-worker-token", request=request),
            WorkerErrorCategory.WORKER_UNAVAILABLE,
        ),
    ):
        calls = 0

        def handler(_: httpx.Request, error: Exception = exception) -> httpx.Response:
            nonlocal calls
            calls += 1
            raise error

        client = worker_client(handler, retries=1)
        with pytest.raises(GenerationWorkerError) as raised:
            await client.health()
        assert raised.value.category is category
        assert "test-worker-token" not in str(raised.value)
        assert calls == 2


@pytest.mark.asyncio
@pytest.mark.parametrize("status_code", [429, 502, 503, 504])
async def test_retryable_http_failures_use_bounded_retry(status_code: int) -> None:
    calls = 0
    delays: list[float] = []

    def handler(_: httpx.Request) -> httpx.Response:
        nonlocal calls
        calls += 1
        if calls < 3:
            return httpx.Response(status_code, json={"secret": "not surfaced"})
        return httpx.Response(200, json={"status": "ok"})

    client = worker_client(handler, retries=2, sleep_calls=delays)
    assert (await client.health()).status is WorkerHealthStatus.HEALTHY
    assert calls == 3
    assert delays == [0, 0]


@pytest.mark.asyncio
@pytest.mark.parametrize("status_code", [400, 401])
async def test_non_retryable_http_failures_do_not_retry(status_code: int) -> None:
    calls = 0

    def handler(_: httpx.Request) -> httpx.Response:
        nonlocal calls
        calls += 1
        return httpx.Response(status_code)

    client = worker_client(handler, retries=3)
    with pytest.raises(GenerationWorkerError) as raised:
        await client.health()
    assert calls == 1
    assert raised.value.http_status == status_code


@pytest.mark.asyncio
async def test_unexpected_exception_is_not_automatically_retryable() -> None:
    calls = 0

    def handler(_: httpx.Request) -> httpx.Response:
        nonlocal calls
        calls += 1
        raise RuntimeError("programming failure with secret-token")

    client = worker_client(handler, retries=3)
    with pytest.raises(GenerationWorkerError) as raised:
        await client.health()
    assert raised.value.category is WorkerErrorCategory.UNKNOWN_ERROR
    assert calls == 1
    assert "secret-token" not in str(raised.value)


@pytest.mark.asyncio
async def test_worker_cancellation_and_output_contract() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        if request.method == "POST":
            return httpx.Response(202, json={"status": "cancellation_requested"})
        return httpx.Response(
            200,
            json={
                "job_id": "job-1",
                "status": "completed",
                "output": {
                    "type": "image",
                    "mime_type": "image/png",
                    "id": "output-1",
                    "download_path": "/v1/outputs/output-1",
                    "filename": "output.png",
                },
            },
        )

    client = worker_client(handler)
    cancellation = await client.cancel("job-1")
    assert cancellation.status is WorkerCancellationStatus.REQUESTED
    output = await client.retrieve_output("job-1")
    assert output.provider_output_id == "output-1"
    assert output.download_path == "/v1/outputs/output-1"
    with pytest.raises(ValidationError):
        WorkerOutput(
            output_type=GenerationModality.IMAGE,
            mime_type="image/png",
            provider_output_id="output-1",
            download_path="https://attacker.example/output.png",
        )
    with pytest.raises(ValidationError):
        WorkerOutput(
            output_type=GenerationModality.IMAGE,
            mime_type="image/png",
            provider_output_id="output-1",
            download_path="/v1/outputs/%2e%2e/secrets",
        )


@pytest.mark.asyncio
async def test_empty_successful_cancellation_response_means_requested_not_cancelled() -> None:
    client = worker_client(lambda _: httpx.Response(204))
    result = await client.cancel("job-1")
    assert result.status is WorkerCancellationStatus.REQUESTED


@pytest.mark.asyncio
async def test_output_stream_is_scoped_to_the_configured_worker_origin() -> None:
    client = worker_client(lambda _: httpx.Response(200, content=b"worker-output"))
    output = WorkerOutput(
        output_type=GenerationModality.IMAGE,
        mime_type="image/png",
        provider_output_id="output-1",
        download_path="/v1/outputs/output-1",
    )
    async with client.stream_output(output) as chunks:
        received = b"".join([chunk async for chunk in chunks])
    assert received == b"worker-output"


@pytest.mark.asyncio
async def test_worker_info_requires_a_discovered_model_match_for_availability() -> None:
    model = GenerationModelCapability(
        id="runtime-image-v1", name="Runtime", modality=GenerationModality.IMAGE
    )
    registry = GenerationModelRegistry(
        [
            GenerationModelRegistration(
                provider_id="runtime-test",
                model=model,
                configuration_reference="runtime-test-config",
            )
        ]
    )
    readiness = WorkerReadiness(
        status=WorkerReadinessStatus.READY, model_loaded=True, model_ids=[model.id]
    )
    undiscovered = WorkerInfo(
        id="worker",
        name="Worker",
        media_types=[GenerationModality.IMAGE],
        models=[{"id": "other-model", "name": "Other", "media_types": ["image"]}],
    )
    assert not registry.verify_readiness(
        provider_id="runtime-test",
        worker_info=undiscovered,
        readiness=readiness,
        provider_configured=True,
    )[0].available


def test_worker_url_and_path_validation_blocks_ssrf_and_traversal() -> None:
    policy = GenerationRetryPolicy(max_retries=0, backoff_seconds=0)
    for url in (
        "http://example.com",
        "https://10.0.0.1",
        "http://169.254.169.254",
        "https://169.254.169.254",
        "https://worker.example/%2e%2e/internal",
        "file:///etc/passwd",
    ):
        with pytest.raises(ValueError):
            RemoteWorkerClient(
                base_url=url,
                bearer_token=None,
                connect_timeout_seconds=1,
                request_timeout_seconds=1,
                read_timeout_seconds=1,
                retry_policy=policy,
            )
    with pytest.raises(GenerationWorkerError):
        RemoteWorkerClient._safe_external_id("job/../../metadata")