From 0cbadec77ebbe866207da69221de7ff2eb457a49 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Wed, 11 Nov 2026 15:14:15 -0800 Subject: [PATCH] fix(apps): include sanitized response details in HTTP errors --- runpod/apps/__init__.py | 2 + runpod/apps/api.py | 6 +- runpod/apps/errors.py | 27 ++- runpod/apps/logs.py | 5 +- runpod/apps/targets.py | 33 ++-- runpod/apps/tasks.py | 5 +- runpod/apps/utils/http.py | 82 +++++++++ runpod/apps/utils/network.py | 3 +- tests/test_apps/test_http_errors.py | 259 ++++++++++++++++++++++++++++ tests/test_apps/test_logs.py | 20 ++- tests/test_apps/test_retry.py | 13 +- tests/test_apps/test_targets.py | 4 +- tests/test_apps/test_tasks.py | 6 +- 13 files changed, 417 insertions(+), 48 deletions(-) create mode 100644 runpod/apps/utils/http.py create mode 100644 tests/test_apps/test_http_errors.py diff --git a/runpod/apps/__init__.py b/runpod/apps/__init__.py index deb993bbf..614233707 100644 --- a/runpod/apps/__init__.py +++ b/runpod/apps/__init__.py @@ -6,6 +6,7 @@ from .entrypoint import local_entrypoint from .errors import ( AppError, + AppHTTPError, EndpointNotFound, JobStatusError, RemoteExecutionError, @@ -26,6 +27,7 @@ "ApiHandle", "App", "AppError", + "AppHTTPError", "Context", "DataCenter", "EndpointNotFound", diff --git a/runpod/apps/api.py b/runpod/apps/api.py index 49dd1d6ca..08ec22f1d 100644 --- a/runpod/apps/api.py +++ b/runpod/apps/api.py @@ -17,6 +17,8 @@ from ..api.rest import run_rest_paginated_request_async, run_rest_request_async from ..error import AuthenticationError, QueryError from ..user_agent import APPS_USER_AGENT +from .errors import AppHTTPError +from .utils.http import raise_for_status _TRANSPORT_RETRIES = 4 @@ -713,7 +715,7 @@ async def _reader(): try: await self._put_tarball(upload_url, _reader(), total, timeout) return - except (aiohttp.ClientError, OSError): + except (AppHTTPError, aiohttp.ClientError, OSError): if attempt == attempts - 1: raise await asyncio.sleep(2**attempt) @@ -729,4 +731,4 @@ async def _put_tarball(self, upload_url, reader, total, timeout) -> None: "Content-Length": str(total), }, ) as resp: - resp.raise_for_status() + await raise_for_status(resp) diff --git a/runpod/apps/errors.py b/runpod/apps/errors.py index 2d34177f0..39a177b5e 100644 --- a/runpod/apps/errors.py +++ b/runpod/apps/errors.py @@ -7,6 +7,18 @@ class AppError(Exception): """base error for the apps sdk.""" +class AppHTTPError(AppError): + """an http failure with a sanitized request url and response excerpt.""" + + def __init__(self, method: str, url: str, status: int, body_excerpt: str = ""): + self.method = method + self.url = str(URL(url).with_user(None).with_query(None).with_fragment(None)) + self.status = status + self.body_excerpt = body_excerpt + detail = f": {body_excerpt}" if body_excerpt else "" + super().__init__(f"HTTP {status} ({method} {self.url}){detail}") + + class EndpointNotFound(AppError): """a resource was invoked but is not deployed.""" @@ -23,18 +35,15 @@ class RemoteExecutionError(AppError): """the remote worker reported a failure executing the function.""" -class JobStatusError(AppError): +class JobStatusError(AppHTTPError): """a job-status request failed with an http error.""" - def __init__(self, job_id: str, status: int, method: str, url: str): + def __init__( + self, job_id: str, status: int, method: str, url: str, body_excerpt: str = "" + ): self.job_id = job_id - self.status = status - self.method = method - self.url = str(URL(url).with_user(None).with_query(None).with_fragment(None)) - super().__init__( - f"could not retrieve status for job '{job_id}': " - f"HTTP {status} ({method} {self.url})" - ) + super().__init__(method, url, status, body_excerpt) + self.args = (f"could not retrieve status for job '{job_id}': {self.args[0]}",) class ScheduleNotSupported(AppError): diff --git a/runpod/apps/logs.py b/runpod/apps/logs.py index 11489fb75..3c8d32f41 100644 --- a/runpod/apps/logs.py +++ b/runpod/apps/logs.py @@ -14,6 +14,7 @@ from ..api.ctl_commands import _path_segment from ..api.rest import _build_url from ..user_agent import APPS_USER_AGENT +from .utils.http import raise_for_status HAPI_BASE = os.environ.get("RUNPOD_HAPI_URL", "https://hapi.runpod.net") @@ -39,7 +40,7 @@ async def pod_logs( async with session.get( url, params={"type": log_type}, headers=_headers() ) as resp: - resp.raise_for_status() + await raise_for_status(resp) data = await resp.json() return {k: v or [] for k, v in data.items()} @@ -88,7 +89,7 @@ async def stream_pod_logs( client_timeout = aiohttp.ClientTimeout(total=STREAM_TIMEOUT_SECONDS, sock_read=60) async with aiohttp.ClientSession(timeout=client_timeout) as session: async with session.get(url, params=params, headers=headers) as resp: - resp.raise_for_status() + await raise_for_status(resp) event_id = None data = [] content = resp.content.__aiter__() diff --git a/runpod/apps/targets.py b/runpod/apps/targets.py index 484857049..97d941243 100644 --- a/runpod/apps/targets.py +++ b/runpod/apps/targets.py @@ -15,11 +15,12 @@ from .._hosts import endpoint_domain from ..user_agent import APPS_USER_AGENT -from .errors import EndpointNotFound, JobStatusError, RemoteExecutionError +from .errors import AppHTTPError, EndpointNotFound, RemoteExecutionError from .protocol import FORMAT_JSON, FunctionRequest from .serialization import require_json_value from .spec import ResourceSpec from .utils.events import emit +from .utils.http import raise_for_status from .utils.network import api_key as _api_key from .utils.network import endpoint_url_base as _endpoint_url_base @@ -246,30 +247,20 @@ async def _request_json( async with session.request( method, url, json=payload, headers=headers ) as resp: - if resp.status >= 400 and job_id is not None: - last_exc = JobStatusError(job_id, resp.status, method, url) - if (resp.status == 404 and retry_not_found) or ( - resp.status in RETRYABLE_STATUSES and retry_transient - ): - continue - raise last_exc - if resp.status == 404: + if resp.status == 404 and job_id is None: last_exc = EndpointNotFound(app_name, resource_name) if retry_not_found and attempt < attempts - 1: continue raise last_exc - if resp.status in RETRYABLE_STATUSES: - last_exc = aiohttp.ClientResponseError( - resp.request_info, - resp.history, - status=resp.status, - message=await resp.text(), - ) - if not retry_transient: - raise last_exc - continue - resp.raise_for_status() + await raise_for_status(resp, job_id=job_id) return await resp.json() + except AppHTTPError as exc: + if (exc.status == 404 and retry_not_found) or ( + exc.status in RETRYABLE_STATUSES and retry_transient + ): + last_exc = exc + continue + raise except (aiohttp.ClientConnectionError, asyncio.TimeoutError) as exc: if not retry_transient: raise @@ -516,7 +507,7 @@ async def http( ) as resp: if resp.status == 404 and self._app_name: raise EndpointNotFound(self._app_name, self._resource_name) - resp.raise_for_status() + await raise_for_status(resp) text = await resp.text() try: result = json.loads(text) diff --git a/runpod/apps/tasks.py b/runpod/apps/tasks.py index 08d4bd0cb..0f4a4bbbc 100644 --- a/runpod/apps/tasks.py +++ b/runpod/apps/tasks.py @@ -24,6 +24,7 @@ from .api import AppsApiClient, is_capacity_error from .context import CLEANUP_TIMEOUT from .spec import ResourceSpec +from .utils.http import raise_for_status log = logging.getLogger(__name__) @@ -285,7 +286,7 @@ async def _post(self, path: str, request: Dict[str, Any]) -> None: if resp.status == 404 and attempt < attempts - 1: await asyncio.sleep(2 * (attempt + 1)) continue - resp.raise_for_status() + await raise_for_status(resp) if timeout is not None: response = await resp.json() if ( @@ -319,7 +320,7 @@ async def poll_result(self) -> Optional[Dict[str, Any]]: if resp.status == 404 and attempt < attempts - 1: await asyncio.sleep(2 * (attempt + 1)) continue - resp.raise_for_status() + await raise_for_status(resp) data = await resp.json() break if data.get("status") == "DONE": diff --git a/runpod/apps/utils/http.py b/runpod/apps/utils/http.py new file mode 100644 index 000000000..3d4948391 --- /dev/null +++ b/runpod/apps/utils/http.py @@ -0,0 +1,82 @@ +"""sanitized diagnostics for apps http responses.""" + +import asyncio +import html +import json +from urllib.parse import quote, quote_plus + +import aiohttp +from yarl import URL + +from ..errors import AppHTTPError, JobStatusError + +BODY_READ_LIMIT = 65536 +BODY_EXCERPT_LIMIT = 1024 + + +async def raise_for_status( + response: aiohttp.ClientResponse, *, job_id: str | None = None +) -> None: + if response.status < 400: + return + + secrets = set() + for item in (*response.history, response): + url = item.request_info.real_url + secrets.update(value for value in url.query.values() if value) + secrets.update(value for value in (url.user, url.password) if value) + for name, value in item.request_info.headers.items(): + if any( + word in name.lower() + for word in ("authorization", "cookie", "key", "token", "secret") + ): + secrets.add(value) + if name.lower() in ("authorization", "proxy-authorization"): + secrets.add(value.partition(" ")[2]) + if name.lower() == "cookie": + secrets.update( + part.partition("=")[2].strip() for part in value.split(";") + ) + + sensitive = set() + for value in secrets: + if value: + sensitive.update( + ( + value, + quote(value, safe=""), + quote_plus(value), + html.escape(value), + json.dumps(value)[1:-1], + ) + ) + + try: + try: + data = await response.content.readexactly(BODY_READ_LIMIT + 1) + except asyncio.IncompleteReadError as exc: + data = exc.partial + body = data.decode("utf-8", errors="replace") + except (aiohttp.ClientError, asyncio.TimeoutError): + body = "[response body unavailable]" + + truncated = len(body) > BODY_EXCERPT_LIMIT + body = body[:BODY_EXCERPT_LIMIT] + url = str(URL(response.url).with_user(None).with_query(None).with_fragment(None)) + for value in sorted(sensitive, key=len, reverse=True): + url = url.replace(value, "[redacted]") + body = body.replace(value, "[redacted]") + if truncated: + for length in range(min(len(value) - 1, len(body)), 0, -1): + if body.endswith(value[:length]): + body = body[:-length] + "[redacted]" + break + body = body[:BODY_EXCERPT_LIMIT] + if truncated: + body = body[: BODY_EXCERPT_LIMIT - 3] + "..." + + if job_id is not None: + raise JobStatusError( + job_id, response.status, response.method, url, body + ) from None + raise AppHTTPError(response.method, url, response.status, body) from None diff --git a/runpod/apps/utils/network.py b/runpod/apps/utils/network.py index af4316fbc..33f4f7588 100644 --- a/runpod/apps/utils/network.py +++ b/runpod/apps/utils/network.py @@ -6,6 +6,7 @@ import aiohttp from ...user_agent import APPS_USER_AGENT +from .http import raise_for_status def runtime_headers() -> Optional[Dict[str, str]]: @@ -65,5 +66,5 @@ async def get_json(url: str, headers: Dict[str, str], timeout: float) -> Dict[st async with session.get( url, headers={**headers, "User-Agent": APPS_USER_AGENT} ) as response: - response.raise_for_status() + await raise_for_status(response) return await response.json() diff --git a/tests/test_apps/test_http_errors.py b/tests/test_apps/test_http_errors.py new file mode 100644 index 000000000..c3de65ea8 --- /dev/null +++ b/tests/test_apps/test_http_errors.py @@ -0,0 +1,259 @@ +"""http diagnostics across the apps transports.""" + +import asyncio +import html +import json +import traceback +from types import SimpleNamespace +from unittest.mock import AsyncMock +from urllib.parse import quote + +import aiohttp +import pytest +from aiohttp import web +from multidict import CIMultiDict +from yarl import URL + +from runpod.apps import AppError, AppHTTPError, JobStatusError +from runpod.apps.api import AppsApiClient +from runpod.apps.logs import pod_logs, stream_pod_logs +from runpod.apps.targets import QueueClient +from runpod.apps.tasks import _TaskTransport +from runpod.apps.utils.http import BODY_EXCERPT_LIMIT, BODY_READ_LIMIT, raise_for_status +from runpod.apps.utils.network import get_json + + +@pytest.fixture +async def error_server(monkeypatch): + state = {"requests": [], "status": 500} + + async def handler(request): + await request.read() + state["requests"].append((request.method, request.path)) + body = ( + "GPU capacity exhausted. " + + request.headers.get("Authorization", "") + + " " + + request.query.get("signature", "") + ) + return web.Response(status=state["status"], text=body) + + app = web.Application() + app.router.add_route("*", "/{path:.*}", handler) + runner = web.AppRunner(app) + await runner.setup() + site = web.TCPSite(runner, "127.0.0.1", 0) + await site.start() + port = site._server.sockets[0].getsockname()[1] + request = aiohttp.ClientSession._request + + async def local_request(self, method, url, **kwargs): + local_url = URL(url).with_scheme("http").with_host("127.0.0.1").with_port(port) + return await request(self, method, local_url, **kwargs) + + monkeypatch.setattr(aiohttp.ClientSession, "_request", local_request) + monkeypatch.setattr("runpod.apps.targets.asyncio.sleep", AsyncMock()) + monkeypatch.setenv("RUNPOD_API_KEY", "synthetic-account-key") + try: + yield state + finally: + await runner.cleanup() + + +@pytest.mark.parametrize( + "operation,method,path,attempts", + [ + ("queue_submit", "POST", "/v2/ep/run", 1), + ("queue_status", "GET", "/v2/ep/status/job", 7), + ("api", "POST", "/predict", 1), + ("task_submit", "POST", "/submit", 1), + ("task_result", "GET", "/result", 1), + ("task_timeout", "POST", "/timeout", 1), + ("upload", "PUT", "/build", 4), + ("snapshot", "GET", "/v1/pod/pod/logs", 1), + ("stream", "GET", "/v1/pod/pod/logs", 1), + ("json", "GET", "/data", 1), + ], +) +async def test_transports_report_body_without_credentials( + error_server, monkeypatch, tmp_path, operation, method, path, attempts +): + monkeypatch.setattr("runpod.endpoint_url_base", "https://api.runpod.ai/v2") + monkeypatch.setattr("runpod.apps.logs.HAPI_BASE", "https://logs.example") + client = QueueClient( + "ep", lambda: {"Authorization": "Bearer synthetic-account-key"} + ) + task = _TaskTransport(None) + task.pod_id = "pod" + task.token = "synthetic-task-token" + with pytest.raises(AppHTTPError) as raised: + if operation == "queue_submit": + await client.run({"input": {}}) + elif operation == "queue_status": + await client.status("job") + elif operation == "api": + await client.http("POST", "/predict", {}) + elif operation == "task_submit": + await task.submit({}) + elif operation == "task_result": + await task.poll_result() + elif operation == "task_timeout": + await task.set_timeout(60) + elif operation == "upload": + tarball = tmp_path / "build.tar.gz" + tarball.write_bytes(b"build") + upload_url = "https://upload.example/build?signature=synthetic-signature" + await AppsApiClient().upload_tarball(upload_url, str(tarball)) + elif operation == "snapshot": + await pod_logs("pod") + elif operation == "stream": + await anext(stream_pod_logs("pod")) + else: + await get_json( + "https://api.example/data", + {"Authorization": "Bearer synthetic-account-key"}, + 10, + ) + + error = raised.value + assert isinstance(error, AppError) + assert error.status == 500 + assert error.method == method + assert error.url.endswith(path) + assert "GPU capacity exhausted." in error.body_excerpt + assert error_server["requests"] == [(method, path)] * attempts + if operation == "queue_status": + assert isinstance(error, JobStatusError) + assert error.job_id == "job" + diagnostic = ( + str(error) + + repr(error) + + repr(vars(error)) + + "".join(traceback.format_exception(type(error), error, error.__traceback__)) + ) + for value in ( + "synthetic-account-key", + "synthetic-task-token", + "synthetic-signature", + ): + assert value not in diagnostic + assert not hasattr(error, "request_info") + assert not hasattr(error, "headers") + assert error.__cause__ is None + assert error.__context__ is None + + +def response_for(body, *, status=500, url="https://api.example/work", headers=None): + url = URL(url) + content = AsyncMock() + if len(body) > BODY_READ_LIMIT: + content.readexactly.return_value = body[: BODY_READ_LIMIT + 1] + else: + content.readexactly.side_effect = asyncio.IncompleteReadError( + body, BODY_READ_LIMIT + 1 + ) + return SimpleNamespace( + status=status, + method="POST", + url=url, + request_info=SimpleNamespace(real_url=url, headers=CIMultiDict(headers or {})), + history=(), + content=content, + ) + + +@pytest.mark.parametrize( + "body,excerpt", + [ + (b'{"error":"quota exceeded"}', '{"error":"quota exceeded"}'), + (b"upstream unavailable", "upstream unavailable"), + (b"", ""), + (b"invalid: \xff", "invalid: \ufffd"), + (b"x" * (BODY_READ_LIMIT * 2), "x" * (BODY_EXCERPT_LIMIT - 3) + "..."), + ], +) +async def test_body_excerpt_is_bounded_and_tolerates_non_json(body, excerpt): + response = response_for(body) + with pytest.raises(AppHTTPError) as raised: + await raise_for_status(response) + assert raised.value.body_excerpt == excerpt + response.content.readexactly.assert_awaited_once_with(BODY_READ_LIMIT + 1) + + +@pytest.mark.parametrize("status", [200, 204, 302]) +async def test_success_does_not_consume_body(status): + response = response_for(b"success", status=status) + await raise_for_status(response) + response.content.readexactly.assert_not_awaited() + + +@pytest.mark.parametrize( + "failure", [aiohttp.ClientPayloadError("secret"), asyncio.TimeoutError()] +) +async def test_unreadable_error_body_preserves_status(failure): + response = response_for(b"") + response.content.readexactly.side_effect = failure + with pytest.raises(AppHTTPError) as raised: + await raise_for_status(response) + assert raised.value.status == 500 + assert raised.value.body_excerpt == "[response body unavailable]" + assert raised.value.__context__ is None + assert raised.value.__cause__ is None + + +async def test_cancellation_during_body_read_propagates(): + response = response_for(b"") + response.content.readexactly.side_effect = asyncio.CancelledError() + with pytest.raises(asyncio.CancelledError): + await raise_for_status(response) + + +async def test_redacts_reflected_credentials_and_redirect_history(): + token = 'synthetic-token<&"' + values = [token, quote(token, safe=""), html.escape(token), json.dumps(token)[1:-1]] + response = response_for( + ( + "rejected " + + " ".join(values) + + " query-secret user-secret password-secret cookie-secret old-token" + ).encode(), + url="https://user-secret:password-secret@api.example/work?key=query-secret#fragment-secret", + headers={"Authorization": f"Bearer {token}", "Cookie": "session=cookie-secret"}, + ) + response.history = ( + response_for(b"", headers={"X-Runpod-Task-Token": "old-token"}), + ) + with pytest.raises(AppHTTPError) as raised: + await raise_for_status(response) + error = raised.value + assert error.url == "https://api.example/work" + assert error.body_excerpt.startswith("rejected [redacted]") + for value in ( + *values, + "query-secret", + "user-secret", + "password-secret", + "cookie-secret", + "old-token", + "fragment-secret", + ): + assert value not in str(error) + repr(error) + repr(vars(error)) + + +async def test_redacts_credentials_crossing_excerpt_boundary(): + token = "synthetic-account-key" + body = ("x" * (BODY_EXCERPT_LIMIT - 10) + token).encode() + response = response_for(body, headers={"Authorization": f"Bearer {token}"}) + with pytest.raises(AppHTTPError) as raised: + await raise_for_status(response) + assert "synthetic" not in raised.value.body_excerpt + assert len(raised.value.body_excerpt) <= BODY_EXCERPT_LIMIT + + +async def test_redacts_complete_repeating_credentials_at_excerpt_boundary(): + token = "synthetic-synthetic-" + body = ("x" * (BODY_EXCERPT_LIMIT - len(token)) + token + "overflow").encode() + response = response_for(body, headers={"Authorization": f"Bearer {token}"}) + with pytest.raises(AppHTTPError) as raised: + await raise_for_status(response) + assert "synthetic" not in raised.value.body_excerpt diff --git a/tests/test_apps/test_logs.py b/tests/test_apps/test_logs.py index 13a5f3a5c..f4f2fa167 100644 --- a/tests/test_apps/test_logs.py +++ b/tests/test_apps/test_logs.py @@ -6,8 +6,11 @@ import pytest from aiohttp import web +from yarl import URL +from runpod.apps import AppHTTPError from runpod.apps.logs import pod_logs, stream_pod_logs, tail_summary +from runpod.apps.utils.http import BODY_READ_LIMIT class TestTailSummary: @@ -35,6 +38,7 @@ def test_system_before_container(self): def _mock_session(response): """a ClientSession whose get() context manager yields `response`.""" + response.status = 200 get_cm = MagicMock() get_cm.__aenter__ = AsyncMock(return_value=response) get_cm.__aexit__ = AsyncMock(return_value=False) @@ -200,9 +204,17 @@ async def test_idle_timeout_starts_after_response_headers( monkeypatch.setenv("RUNPOD_API_KEY", "test-key") response = MagicMock() response.content = _FakeContent([b'data: {"line": "ready"}\n', b"\n"]) - if status == 403: - response.raise_for_status.side_effect = web.HTTPForbidden() session_cm, session = _mock_session(response) + response.status = status + response.method = "GET" + response.url = response.request_info.real_url = URL("https://logs.example/pod1") + if status == 403: + response.content = MagicMock() + response.content.readexactly = AsyncMock( + side_effect=asyncio.IncompleteReadError( + b"forbidden", BODY_READ_LIMIT + 1 + ) + ) async def connect(): await asyncio.sleep(0.03) @@ -212,8 +224,10 @@ async def connect(): with patch("aiohttp.ClientSession", return_value=session_cm): stream = stream_pod_logs("pod1", idle_timeout=0.001) if status == 403: - with pytest.raises(web.HTTPForbidden): + with pytest.raises(AppHTTPError) as raised: await anext(stream) + assert raised.value.status == 403 + assert raised.value.body_excerpt == "forbidden" else: assert [e async for e in stream] == [{"line": "ready"}] session.get.return_value.__aexit__.assert_awaited_once() diff --git a/tests/test_apps/test_retry.py b/tests/test_apps/test_retry.py index abd732d5f..09c608be3 100644 --- a/tests/test_apps/test_retry.py +++ b/tests/test_apps/test_retry.py @@ -10,7 +10,7 @@ import aiohttp import pytest -from runpod.apps.errors import EndpointNotFound +from runpod.apps.errors import AppHTTPError, EndpointNotFound from runpod.apps.targets import ( JOB_PROPAGATION_ATTEMPTS, RETRY_ATTEMPTS, @@ -77,9 +77,10 @@ async def test_retries_transient_5xx_then_succeeds(flaky_server): async def test_exhausted_retries_raise_last_error(flaky_server): flaky_server.state["script"] = [503] * (RETRY_ATTEMPTS + 2) with patch("runpod.apps.targets.RETRY_BASE_DELAY", 0.01): - with pytest.raises(aiohttp.ClientResponseError) as exc_info: + with pytest.raises(AppHTTPError) as exc_info: await _request_json("GET", flaky_server.url, {}, timeout=10) assert exc_info.value.status == 503 + assert exc_info.value.body_excerpt == "scripted" assert flaky_server.state["hits"] == RETRY_ATTEMPTS @@ -127,9 +128,10 @@ async def test_persistent_job_404_exhausts_retries(flaky_server): async def test_client_4xx_not_retried(flaky_server): flaky_server.state["script"] = [400] - with pytest.raises(aiohttp.ClientResponseError) as exc_info: + with pytest.raises(AppHTTPError) as exc_info: await _post_json(flaky_server.url, {"x": 1}, {}, timeout=10) assert exc_info.value.status == 400 + assert exc_info.value.body_excerpt == "scripted" assert flaky_server.state["hits"] == 1 @@ -154,7 +156,10 @@ async def test_submission_is_not_repeated_after_ambiguous_failure( else: flaky_server.state["script"] = [failure] client = QueueClient("endpoint", lambda: {}) - with pytest.raises(aiohttp.ClientError): + error_type = ( + aiohttp.ClientConnectionError if failure == "disconnect" else AppHTTPError + ) + with pytest.raises(error_type): if operation == "invoke": await client.invoke({"input": {}}, timeout=10) elif operation == "retry": diff --git a/tests/test_apps/test_targets.py b/tests/test_apps/test_targets.py index dfe2cbd69..84e46c461 100644 --- a/tests/test_apps/test_targets.py +++ b/tests/test_apps/test_targets.py @@ -545,14 +545,14 @@ async def test_concurrent_requests_share_completed_source_sync( async def test_failed_source_sync_prevents_execution_and_can_retry( self, live_api_server ): - import aiohttp + from runpod.apps import AppHTTPError state, _, release = live_api_server state["reject_sync"] = True release.set() target = LiveTarget("ep123", "calculate") target.attach_source(self._spec, "calculate", self._spec()) - with pytest.raises(aiohttp.ClientResponseError): + with pytest.raises(AppHTTPError): await target.request("POST", "/calculate", {"value": 7}) assert state["calls"] == 0 assert await target.request("POST", "/calculate", {"value": 7}) == { diff --git a/tests/test_apps/test_tasks.py b/tests/test_apps/test_tasks.py index cab2b122f..ab2c4d88d 100644 --- a/tests/test_apps/test_tasks.py +++ b/tests/test_apps/test_tasks.py @@ -275,13 +275,15 @@ def test_retries_transient_404_then_returns_result(self): httpd.server_close() def test_persistent_404_raises(self): - import aiohttp + from runpod.apps import AppHTTPError httpd, hits = self._server([(404, {"error": "not found"})]) try: url = f"http://127.0.0.1:{httpd.server_port}" - with pytest.raises(aiohttp.ClientResponseError): + with pytest.raises(AppHTTPError) as raised: self._poll(url) + assert raised.value.status == 404 + assert raised.value.body_excerpt == '{"error": "not found"}' assert hits["count"] == 6 finally: httpd.shutdown()