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