From 2448bd4f12a750b9410fdd8fbcd57bfcd2ba5c63 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 1 Oct 2026 19:10:40 -0700 Subject: [PATCH 1/7] add async authorized session support --- .../resumable_transfer/upload_async.py | 270 +++++++++----- .../asyncio/test_resumable_transfer_async.py | 350 +++++++++++++++++- 2 files changed, 519 insertions(+), 101 deletions(-) diff --git a/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py b/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py index 13cfcb5c6fbd..9dab3b749f9a 100644 --- a/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py +++ b/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Asynchronous Resumable Upload session and helpers using aiohttp.""" +"""Asynchronous Resumable Upload session and helpers.""" import asyncio import datetime @@ -28,6 +28,7 @@ Awaitable, BinaryIO, Callable, + Dict, Generator, Generic, Iterable, @@ -35,8 +36,10 @@ Mapping, Optional, Tuple, + Type, TypeVar, Union, + cast, ) try: @@ -46,6 +49,14 @@ except ImportError: # pragma: NO COVER _HAS_AIOHTTP = False +try: + from google.auth import exceptions as auth_exceptions + from google.auth.aio.transport.sessions import AsyncAuthorizedSession + + _HAS_GOOGLE_AUTH_AIO = True +except ImportError: # pragma: NO COVER + _HAS_GOOGLE_AUTH_AIO = False + import google.api_core.retry from google.api_core import exceptions from google.api_core.resumable_transfer import common, upload_state @@ -60,6 +71,34 @@ _monotonic_clock = time.monotonic +AsyncTransport = Union["aiohttp.ClientSession", "AsyncAuthorizedSession"] +"""Asynchronous HTTP transports accepted by :class:`AsyncResumableUploadSession`. + +``google.auth.aio.transport.sessions.AsyncAuthorizedSession`` is the transport +used by generated asynchronous REST clients. +""" + +# Timeouts and transport-level failures raised by the supported transports. +# ``AsyncAuthorizedSession`` surfaces aiohttp errors as google.auth exceptions. +_TIMEOUT_ERRORS: Tuple[Type[BaseException], ...] = (asyncio.TimeoutError,) +_TRANSPORT_ERRORS: Tuple[Type[BaseException], ...] = () +# Extra keyword arguments for every ``AsyncAuthorizedSession.request()`` call. +_AUTHORIZED_SESSION_REQUEST_KWARGS: Dict[str, Any] = {} +if _HAS_AIOHTTP: # pragma: no branch + _TIMEOUT_ERRORS += (aiohttp.ServerTimeoutError,) + _TRANSPORT_ERRORS += (aiohttp.ClientError,) +if _HAS_GOOGLE_AUTH_AIO: # pragma: no branch + _TIMEOUT_ERRORS += (auth_exceptions.TimeoutError,) + _TRANSPORT_ERRORS += (auth_exceptions.TransportError, auth_exceptions.ResponseError) + # google-auth >= 2.49.0 accepts ``total_attempts``; a single attempt disables + # its internal status-code retries so this module is the only retry layer. + _request_params = inspect.signature(AsyncAuthorizedSession.request).parameters + if "total_attempts" in _request_params: # pragma: no branch + _AUTHORIZED_SESSION_REQUEST_KWARGS["total_attempts"] = 1 +# Errors that may carry a chunk timeout out of the streaming retry loop. +_TIMEOUT_OR_RETRY_ERRORS = (exceptions.RetryError,) + _TIMEOUT_ERRORS + + def _get_buffer_size(stream: object) -> Optional[int]: """Returns buffer size in bytes if stream exposes getbuffer(), else None.""" getbuffer_fn = getattr(stream, "getbuffer", None) @@ -166,7 +205,7 @@ def __init__( self, upload_url: Optional[str] = None, config: Optional[ResumableUploadConfig] = None, - transport: Optional[Any] = None, + transport: Optional[AsyncTransport] = None, content_type: Optional[str] = None, response_type: Optional[Any] = None, start_retry: Optional[google.api_core.retry.AsyncRetry] = None, @@ -179,7 +218,7 @@ def __init__( new upload, or the pre-existing upload session URL when resuming. config: Optional upload configuration parameters. Defaults to ``ResumableUploadConfig()`` when ``None``. - transport: Optional aiohttp.ClientSession. When ``None``, a + transport: Optional :data:`AsyncTransport`. When ``None``, a transport must be provided to ``upload()`` or ``resume()``. content_type: Optional MIME type of the stream payload. When ``None``, no content-type header is sent unless overridden. @@ -276,6 +315,80 @@ def _ensure_aiohttp(self) -> None: "Please install google-api-core[async_rest]." ) + def _get_transport(self, transport: Optional[AsyncTransport]) -> AsyncTransport: + """Returns the transport to use, raising ValueError when none is set.""" + sess = transport or self._transport + if sess is None: + raise ValueError( + "An aiohttp.ClientSession or AsyncAuthorizedSession transport " + "must be provided." + ) + return sess + + async def _send_request( + self, + transport: AsyncTransport, + method: str, + url: str, + payload: bytes, + headers: Mapping[str, str], + timeout: float, + ) -> Tuple[int, Mapping[str, str], bytes]: + """Sends one request over ``transport`` and reads the whole response. + + Args: + transport: The transport to dispatch the request over. + method: HTTP method verb. + url: Request URL. + payload: Request body bytes. + headers: Request headers. + timeout: Total timeout in seconds for this attempt. + + Returns: + Tuple of (status code, response headers, response body bytes). + + Raises: + exceptions.GoogleAPICallError: If the status is not 200 or 201. + """ + if _HAS_GOOGLE_AUTH_AIO and isinstance(transport, AsyncAuthorizedSession): + # ``request()`` is a coroutine returning a ``Response`` that must be + # closed explicitly. ``max_allowed_time`` keeps google-auth's + # wall-clock guard within the same per-attempt timeout. + response = await transport.request( + method, + url, + data=payload, + headers=headers, + timeout=timeout, + max_allowed_time=timeout, + **_AUTHORIZED_SESSION_REQUEST_KWARGS, + ) + try: + status_code = response.status_code + resp_headers: Mapping[str, str] = dict(response.headers) + body = await response.read() + finally: + await response.close() + else: + # mypy cannot narrow the negative of the guarded check above. + session = cast("aiohttp.ClientSession", transport) + async with session.request( + method, + url, + data=payload, + headers=headers, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as resp: + status_code = resp.status + resp_headers = dict(resp.headers) + body = await resp.read() + + if status_code not in (200, 201): + raise exceptions.from_http_status( + status_code, body.decode("utf-8", errors="replace") + ) + return status_code, resp_headers, body + def _enrich_exception(self, exc: BaseException) -> None: """Attaches session diagnostic metadata to an active exception. @@ -366,8 +479,9 @@ def _get_retry_predicate( ``UnseekableStreamError``) always return ``False``. 2. Protocol-recoverable errors during chunk transfer (``RECOVERABLE_STATUS_CODES`` and ``MissingStatusHeaderError``) - and transport errors (``aiohttp.ClientError`` and - ``asyncio.TimeoutError``) always return ``True`` so the session + and transport errors (``aiohttp.ClientError``, + ``asyncio.TimeoutError``, and their ``google.auth.exceptions`` + counterparts) always return ``True`` so the session can query server state and recover. Returns: @@ -385,9 +499,7 @@ def should_retry(exc: Exception) -> bool: ) ): return True - if isinstance(exc, asyncio.TimeoutError) or ( - _HAS_AIOHTTP and isinstance(exc, aiohttp.ClientError) - ): + if isinstance(exc, _TIMEOUT_ERRORS + _TRANSPORT_ERRORS): return True if ( custom_predicate is not None @@ -453,7 +565,7 @@ def _get_async_streaming_retry( async def _initiate( self, - transport: Any, + transport: AsyncTransport, request_body: Union[str, bytes] = "", size: Optional[int] = None, progress_queue: Optional[List[UploadProgress]] = None, @@ -462,7 +574,7 @@ async def _initiate( """Initiates the upload session asynchronously. Args: - transport: The aiohttp client session. + transport: The asynchronous transport to dispatch requests over. request_body: Initial metadata payload sent with start request. size: Total stream size in bytes, if known. progress_queue: Optional queue to receive progress event. @@ -483,21 +595,15 @@ async def _initiate( ) async def do_initiate() -> str: - timeout_sec = self._get_start_timeout() - client_timeout = aiohttp.ClientTimeout(total=timeout_sec) - async with transport.request( - method, url, data=payload, headers=headers, timeout=client_timeout - ) as resp: - resp_headers = dict(resp.headers) - body = await resp.read() - if resp.status not in (200, 201): - raise exceptions.from_http_status( - resp.status, body.decode("utf-8", errors="replace") - ) - session_url = self._state.process_start_response( - resp.status, resp_headers - ) - return session_url + status_code, resp_headers, _ = await self._send_request( + transport, + method, + url, + payload, + headers, + timeout=self._get_start_timeout(), + ) + return self._state.process_start_response(status_code, resp_headers) retry_policy = self._get_async_retry() retryable_initiate = retry_policy(do_initiate) @@ -507,7 +613,7 @@ async def do_initiate() -> str: async def _transmit_chunk( self, - transport: Any, + transport: AsyncTransport, reader_fn: Callable[[int], Awaitable[bytes]], size: Optional[int], progress_queue: Optional[List[UploadProgress]] = None, @@ -516,7 +622,7 @@ async def _transmit_chunk( """Transmits the next data chunk asynchronously with stall control. Args: - transport: The aiohttp client session. + transport: The asynchronous transport to dispatch requests over. reader_fn: Async callable returning chunk bytes. size: Total stream size in bytes, if known. progress_queue: Optional queue to receive progress updates. @@ -563,33 +669,20 @@ async def _transmit_chunk( per_attempt_timeout = self._compute_chunk_timeout( data_len, timeout_override=timeout ) - client_timeout = aiohttp.ClientTimeout(total=per_attempt_timeout) t_start = _monotonic_clock() try: - async with transport.request( + status_code, resp_headers, resp_body = await self._send_request( + transport, method, url, - data=payload, - headers=headers, - timeout=client_timeout, - ) as resp: - resp_headers = dict(resp.headers) - resp_body = await resp.read() - if resp.status not in (200, 201): - raise exceptions.from_http_status( - resp.status, resp_body.decode("utf-8", errors="replace") - ) - status_code = resp.status + payload, + headers, + timeout=per_attempt_timeout, + ) t_elapsed = _monotonic_clock() - t_start except Exception as exc: self._enrich_exception(exc) - if isinstance( - exc, - ( - asyncio.TimeoutError, - aiohttp.ServerTimeoutError, - ), - ): + if isinstance(exc, _TIMEOUT_ERRORS): t_elapsed = _monotonic_clock() - t_start self._get_deadline_remaining() self._update_stall_control(0, t_elapsed) @@ -684,14 +777,14 @@ def _update_stall_control(self, data_len: int, t_elapsed: float) -> None: async def _recover( self, - transport: Any, + transport: AsyncTransport, stream_obj: Optional[object] = None, progress_queue: Optional[List[UploadProgress]] = None, ) -> Tuple[int, Mapping[str, str], bytes]: """Queries server for committed byte offset and adjusts buffer. Args: - transport: The aiohttp client session. + transport: The asynchronous transport to dispatch requests over. stream_obj: Underlying stream object to rewind if seekable. progress_queue: Optional queue to receive progress updates. @@ -703,18 +796,14 @@ async def _recover( exceptions.GoogleAPICallError: If query request fails on the server. """ method, url, headers, payload = self._state.build_query_request() - timeout_sec = self._get_start_timeout() - client_timeout = aiohttp.ClientTimeout(total=timeout_sec) - async with transport.request( - method, url, data=payload, headers=headers, timeout=client_timeout - ) as resp: - resp_headers = dict(resp.headers) - body = await resp.read() - if resp.status not in (200, 201): - raise exceptions.from_http_status( - resp.status, body.decode("utf-8", errors="replace") - ) - status_code = resp.status + status_code, resp_headers, body = await self._send_request( + transport, + method, + url, + payload, + headers, + timeout=self._get_start_timeout(), + ) received = self._state.process_query_response(status_code, resp_headers) self._notify_progress(common.ProgressState.OFFSET_RECEIVED, progress_queue) @@ -760,37 +849,32 @@ async def _recover( chunk_size=self.chunk_size, ) - async def cancel(self, transport: Optional[Any] = None) -> None: + async def cancel(self, transport: Optional[AsyncTransport] = None) -> None: """Cancels the resumable upload session asynchronously. Args: - transport: Optional aiohttp client session. + transport: Optional aiohttp.ClientSession or AsyncAuthorizedSession. Raises: ValueError: If transport is missing. exceptions.GoogleAPICallError: If cancellation request fails on the server. """ self._ensure_aiohttp() - sess = transport or self._transport - if sess is None: - raise ValueError("An aiohttp.ClientSession transport must be provided.") + sess = self._get_transport(transport) method, url, headers, payload = self._state.build_cancel_request() - timeout_sec = self._get_start_timeout() - client_timeout = aiohttp.ClientTimeout(total=timeout_sec) - async with sess.request( - method, url, data=payload, headers=headers, timeout=client_timeout - ) as resp: - resp_headers = dict(resp.headers) - body = await resp.read() - if resp.status not in (200, 201): - raise exceptions.from_http_status( - resp.status, body.decode("utf-8", errors="replace") - ) - self._state.process_cancel_response(resp.status, resp_headers) + status_code, resp_headers, _ = await self._send_request( + sess, + method, + url, + payload, + headers, + timeout=self._get_start_timeout(), + ) + self._state.process_cancel_response(status_code, resp_headers) async def _transmit_all_chunks( self, - transport: Any, + transport: AsyncTransport, reader_fn: Callable[[int], Awaitable[bytes]], computed_size: Optional[int], progress_queue: Optional[List[UploadProgress]] = None, @@ -801,7 +885,7 @@ async def _transmit_all_chunks( """Transmits chunks until completion using a single outer AsyncStreamingRetry coordinator. Args: - transport: The aiohttp client session. + transport: The asynchronous transport to dispatch requests over. reader_fn: Async callable returning chunk bytes. computed_size: Total stream size in bytes, if known. progress_queue: Optional list receiving UploadProgress snapshots. @@ -856,17 +940,11 @@ async def attempt_stream() -> AsyncGenerator[UploadProgress, None]: stream_gen = await retryable_stream() async for item in stream_gen: yield item - except ( - asyncio.TimeoutError, - aiohttp.ServerTimeoutError, - exceptions.RetryError, - ) as exc: + except _TIMEOUT_OR_RETRY_ERRORS as exc: timeout_exc = ( exc.__cause__ if isinstance(exc, exceptions.RetryError) else exc ) - if not isinstance( - timeout_exc, (asyncio.TimeoutError, aiohttp.ServerTimeoutError) - ): + if not isinstance(timeout_exc, _TIMEOUT_ERRORS): raise self._enrich_exception(timeout_exc) self._get_deadline_remaining() @@ -887,7 +965,7 @@ def upload( stream: Union[AsyncIterable[bytes], BinaryIO, bytes, Iterable[bytes]], request_body: Union[str, bytes] = "", size: Optional[int] = None, - transport: Optional[Any] = None, + transport: Optional[AsyncTransport] = None, content_type: Optional[str] = None, retry: Optional[google.api_core.retry.AsyncStreamingRetry] = None, timeout: Optional[float] = None, @@ -898,7 +976,7 @@ def upload( stream: Data payload to upload (async iterable, binary stream, bytes, or iterable). request_body: Initial metadata payload sent with the start request. size: Total stream size in bytes, if known. - transport: Optional aiohttp client session. + transport: Optional aiohttp.ClientSession or AsyncAuthorizedSession. content_type: Optional MIME type of the stream payload. retry: Optional retry configuration (``AsyncStreamingRetry``) for chunk upload requests. Use this to customize exponential backoff timing between chunk retries or to @@ -919,9 +997,7 @@ def upload( ValueError: If transport is missing. """ self._ensure_aiohttp() - sess = transport or self._transport - if sess is None: - raise ValueError("An aiohttp.ClientSession transport must be provided.") + sess = self._get_transport(transport) if content_type is not None: self._content_type = content_type @@ -963,7 +1039,7 @@ def resume( stream: Union[AsyncIterable[bytes], BinaryIO, bytes, Iterable[bytes]], size: Optional[int] = None, chunk_size: Optional[int] = None, - transport: Optional[Any] = None, + transport: Optional[AsyncTransport] = None, retry: Optional[google.api_core.retry.AsyncStreamingRetry] = None, timeout: Optional[float] = None, ) -> AsyncUploadOperation: @@ -975,7 +1051,7 @@ def resume( Data payload to resume uploading. size (Optional[int]): Total stream size in bytes, if known. chunk_size (Optional[int]): Optional chunk size override in bytes. - transport (Optional[Any]): Optional aiohttp client session. + transport (Optional[AsyncTransport]): Optional transport override. retry (Optional[google.api_core.retry.AsyncStreamingRetry]): Optional retry configuration (``AsyncStreamingRetry``) for chunk upload requests. Use this to customize exponential backoff timing @@ -998,9 +1074,7 @@ def resume( ValueError: If transport, upload_url, or stream is missing. """ self._ensure_aiohttp() - sess = transport or self._transport - if sess is None: - raise ValueError("An aiohttp.ClientSession transport must be provided.") + sess = self._get_transport(transport) actual_url = upload_url or self.upload_url if not actual_url: raise ValueError("An upload URL must be provided to resume.") diff --git a/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py b/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py index b16211795ee6..3a832b3d030c 100644 --- a/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py +++ b/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py @@ -14,8 +14,14 @@ """Asynchronous tests for Resumable Upload protocol implementation.""" +# The transport doubles below only implement the parts of aiohttp.ClientSession +# that the session uses, so they do not satisfy the nominal AsyncTransport union. +# mypy: disable-error-code="arg-type" + import asyncio +import contextlib import datetime +import inspect import io import json from typing import ( @@ -32,6 +38,7 @@ from unittest import mock import pytest +from google.auth import exceptions as auth_exceptions from google.protobuf import empty_pb2 from google.api_core import exceptions @@ -52,6 +59,11 @@ try: import aiohttp # noqa: F401 import google.auth.aio.transport # noqa: F401 + from aiohttp import web + from aiohttp.test_utils import TestServer + from google.auth.aio import credentials as aio_credentials + from google.auth.aio.transport import aiohttp as auth_aiohttp + from google.auth.aio.transport import sessions as aio_sessions GOOGLE_AUTH_AIO_INSTALLED = True except ImportError: @@ -1325,12 +1337,12 @@ def test_async_transport_missing_errors() -> None: upload_url="https://api.example.com/start", ) with pytest.raises( - ValueError, match="An aiohttp.ClientSession transport must be provided" + ValueError, match="or AsyncAuthorizedSession transport must be provided" ): session.upload(stream=b"data") with pytest.raises( - ValueError, match="An aiohttp.ClientSession transport must be provided" + ValueError, match="or AsyncAuthorizedSession transport must be provided" ): session.resume(upload_url="https://upload.example.com/123", stream=b"data") @@ -1342,7 +1354,7 @@ async def test_async_cancel_missing_transport_and_error() -> None: upload_url="https://upload.example.com/123", ) with pytest.raises( - ValueError, match="An aiohttp.ClientSession transport must be provided" + ValueError, match="or AsyncAuthorizedSession transport must be provided" ): await session.cancel() @@ -2292,3 +2304,335 @@ def test_async_retry_predicate_includes_asyncio_timeout_error() -> None: session = AsyncResumableUploadSession(upload_url="https://api.example.com/start") predicate = session._get_retry_predicate() assert predicate(asyncio.TimeoutError()) is True + + +# ===================================================================== +# 11. google.auth AsyncAuthorizedSession Transport Tests +# ===================================================================== + + +START_HEADERS = { + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/123", +} + + +def _auth_response( + status_code: int = 200, + headers: Optional[Mapping[str, str]] = None, + body: bytes = b"", +) -> mock.Mock: + """Builds a ``google.auth.aio.transport.Response`` double.""" + return mock.Mock( + status_code=status_code, + headers=dict(headers or {}), + read=mock.AsyncMock(return_value=body), + close=mock.AsyncMock(), + ) + + +def _authorized_session( + *outcomes: Union[mock.Mock, BaseException], +) -> Tuple["aio_sessions.AsyncAuthorizedSession", mock.AsyncMock]: + """Builds a real AsyncAuthorizedSession whose HTTP adapter replays ``outcomes``. + + google-auth invokes the adapter positionally as + ``(url, method, body, headers, timeout, **kwargs)``. + """ + auth_request = mock.AsyncMock(side_effect=list(outcomes)) + transport = aio_sessions.AsyncAuthorizedSession( + aio_credentials.AnonymousCredentials(), auth_request=auth_request + ) + return transport, auth_request + + +def _adapter_commands(auth_request: mock.AsyncMock) -> List[str]: + return [c.args[3]["X-Goog-Upload-Command"] for c in auth_request.call_args_list] + + +def _supports_total_attempts() -> bool: + """Returns True if google-auth (>= 2.49.0) accepts ``total_attempts``.""" + params = inspect.signature(aio_sessions.AsyncAuthorizedSession.request).parameters + return "total_attempts" in params + + +@pytest.mark.asyncio +async def test_async_authorized_session_upload() -> None: + """Uploads through a real AsyncAuthorizedSession (the generated REST client transport). + + Unlike aiohttp.ClientSession, its ``request()`` is a coroutine returning a + ``google.auth.aio.transport.Response`` (``status_code``, explicit ``close()``). + """ + responses = [ + _auth_response(200, START_HEADERS), + _auth_response(200, {"X-Goog-Upload-Status": "active"}), + _auth_response( + 200, {"X-Goog-Upload-Status": "final"}, b'{"name": "auth.txt", "size": 6}' + ), + ] + transport, auth_request = _authorized_session(*responses) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=ResumableUploadConfig(chunk_size=4), + transport=transport, + response_type=DummyResponse, + start_timeout=7.5, + ) + + with mock.patch.object(transport, "request", wraps=transport.request) as spy: + result = await session.upload(stream=b"012345") + + assert isinstance(result, DummyResponse) + assert (result.name, result.size) == ("auth.txt", 6) + assert session.bytes_uploaded == 6 + assert _adapter_commands(auth_request) == ["start", "upload", "upload, finalize"] + assert [c.args[2] for c in auth_request.call_args_list] == [b"", b"0123", b"45"] + assert [c.args[0] for c in auth_request.call_args_list] == [ + "https://api.example.com/start", + "https://upload.example.com/123", + "https://upload.example.com/123", + ] + for response in responses: + response.close.assert_awaited_once() + # api_core owns the timeout: google-auth's wall-clock guard is bounded by the + # same per-attempt timeout, and nothing extra leaks through to the adapter. + assert spy.call_args_list[0].kwargs["timeout"] == 7.5 + expected_total_attempts = 1 if _supports_total_attempts() else None + for call in spy.call_args_list: + assert call.kwargs["max_allowed_time"] == call.kwargs["timeout"] + assert call.kwargs.get("total_attempts") == expected_total_attempts + assert all(c.kwargs == {} for c in auth_request.call_args_list) + + +@pytest.mark.asyncio +async def test_async_authorized_session_server_error_not_retried_internally() -> None: + """Verifies a 5xx reaches api_core's retry layer instead of google-auth's. + + AsyncAuthorizedSession retries retryable status codes itself by default; + api_core passes ``total_attempts=1`` (google-auth >= 2.49.0) so that its + protocol-aware retry and offset recovery is the single retry layer. + """ + if not _supports_total_attempts(): + pytest.skip("google-auth < 2.49.0 retries 5xx responses internally") + response = _auth_response(503, body=b"Service Unavailable") + transport, auth_request = _authorized_session(response) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + transport=transport, + start_retry=google.api_core.retry.AsyncRetry(predicate=lambda exc: False), + ) + + with pytest.raises(exceptions.ServiceUnavailable, match="Service Unavailable"): + await session.upload(stream=b"data") + + assert auth_request.await_count == 1 + response.close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_async_authorized_session_timeout_raises_transfer_stalled() -> None: + """Verifies google.auth.exceptions.TimeoutError feeds stall control. + + It is not a subclass of asyncio.TimeoutError, so it must be recognised + explicitly to surface as TransferStalledError. + """ + transport, auth_request = _authorized_session( + _auth_response(200, START_HEADERS), auth_exceptions.TimeoutError("timed out") + ) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", transport=transport + ) + no_retry = google.api_core.retry.AsyncStreamingRetry( + predicate=lambda e: False, timeout=0.001 + ) + + with pytest.raises(TransferStalledError) as exc_info: + await session.upload(stream=b"data", retry=no_retry) + + assert isinstance(exc_info.value.__cause__, auth_exceptions.TimeoutError) + assert _adapter_commands(auth_request) == ["start", "upload, finalize"] + + +@pytest.mark.asyncio +async def test_async_authorized_session_transport_error_recovers() -> None: + """Verifies google.auth.exceptions.TransportError triggers offset recovery.""" + transport, auth_request = _authorized_session( + _auth_response(200, START_HEADERS), + auth_exceptions.TransportError("connection reset"), + _auth_response( + 200, {"X-Goog-Upload-Status": "active", "X-Goog-Upload-Size-Received": "0"} + ), + _auth_response(200, {"X-Goog-Upload-Status": "final"}, b'{"name": "f"}'), + ) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + transport=transport, + response_type=DummyResponse, + ) + + result = await session.upload( + stream=b"data", + retry=google.api_core.retry.AsyncStreamingRetry(initial=0.01, maximum=0.01), + ) + + assert isinstance(result, DummyResponse) + assert result.name == "f" + assert _adapter_commands(auth_request) == [ + "start", + "upload, finalize", + "query", + "upload, finalize", + ] + + +def test_async_retry_predicate_includes_google_auth_errors() -> None: + """Verifies google-auth transport/timeout errors are retryable, unlike its other errors.""" + session = AsyncResumableUploadSession(upload_url="https://api.example.com/start") + for is_start in (True, False): + predicate = session._get_retry_predicate(is_start=is_start) + assert predicate(auth_exceptions.TimeoutError()) is True + assert predicate(auth_exceptions.TransportError()) is True + assert predicate(auth_exceptions.ResponseError()) is True + assert predicate(auth_exceptions.RefreshError()) is False + + +# ===================================================================== +# 12. End-to-End Tests Against a Local HTTP Server +# ===================================================================== + + +class ResumableUploadServer: + """Minimal in-memory resumable upload server. + + ``stalls`` maps a zero-based request ordinal to seconds to sleep before + answering that request with 408 without committing any bytes. + """ + + def __init__(self) -> None: + self.blobs: Dict[str, bytearray] = {} + self.commands: List[str] = [] + self.cancelled: List[str] = [] + self.stalls: Dict[int, float] = {} + + def make_app(self) -> "web.Application": + app = web.Application() + app.router.add_post("/start", self.handle) + app.router.add_post("/upload/{blob}", self.handle) + return app + + async def handle(self, request: "web.Request") -> "web.Response": + command = request.headers.get(common.HEADER_COMMAND, "") + payload = await request.read() + ordinal = len(self.commands) + self.commands.append(command) + if ordinal in self.stalls: + await asyncio.sleep(self.stalls.pop(ordinal)) + return web.Response(status=408, text="stalled") + if command == common._Command.START.value: + blob = f"blob-{len(self.blobs)}" + self.blobs[blob] = bytearray() + headers = { + common.HEADER_STATUS: common._Status.ACTIVE.value, + common.HEADER_URL: str(request.url.with_path(f"/upload/{blob}")), + } + return web.Response(headers=headers) + blob = request.match_info["blob"] + data = self.blobs[blob] + if command == common._Command.QUERY.value: + headers = { + common.HEADER_STATUS: common._Status.ACTIVE.value, + common.HEADER_SIZE_RECEIVED: str(len(data)), + } + return web.Response(headers=headers) + if command == common._Command.CANCEL.value: + self.cancelled.append(blob) + headers = {common.HEADER_STATUS: common._Status.CANCELLED.value} + return web.Response(headers=headers) + if int(request.headers.get(common.HEADER_OFFSET, "-1")) != len(data): + return web.Response(status=400, text="offset mismatch") + data.extend(payload) + if common._Command.FINALIZE.value in command: + headers = {common.HEADER_STATUS: common._Status.FINAL.value} + return web.json_response({"name": blob, "size": len(data)}, headers=headers) + return web.Response(headers={common.HEADER_STATUS: common._Status.ACTIVE.value}) + + +@contextlib.asynccontextmanager +async def serve_resumable_upload() -> AsyncIterator[Tuple[ResumableUploadServer, str]]: + """Serves a ``ResumableUploadServer`` locally, yielding it and its base URL.""" + server = ResumableUploadServer() + async with TestServer(server.make_app()) as test_server: + yield server, str(test_server.make_url("")).rstrip("/") + + +@pytest.mark.asyncio +async def test_async_upload_end_to_end_aiohttp_client_session() -> None: + """Uploads over the wire with a real aiohttp.ClientSession.""" + async with serve_resumable_upload() as (server, base_url): + async with aiohttp.ClientSession() as client: + session = AsyncResumableUploadSession( + upload_url=f"{base_url}/start", + config=ResumableUploadConfig(chunk_size=4), + transport=client, + response_type=DummyResponse, + ) + result = await session.upload(stream=b"0123456789") + + assert isinstance(result, DummyResponse) + assert bytes(server.blobs[result.name]) == b"0123456789" + assert server.commands == ["start", "upload", "upload", "upload, finalize"] + + +@pytest.mark.asyncio +async def test_async_upload_end_to_end_authorized_session() -> None: + """Uploads, resumes, and cancels over the wire with the generated client transport. + + A real AsyncAuthorizedSession backed by google-auth's aiohttp adapter is + used. The first chunk request stalls past the per-attempt timeout, so + google-auth raises its own TimeoutError and api_core must query the + committed offset and re-send the chunk. + """ + quick_retry = google.api_core.retry.AsyncStreamingRetry(initial=0.01, maximum=0.01) + async with serve_resumable_upload() as (server, base_url): + transport = aio_sessions.AsyncAuthorizedSession( + aio_credentials.AnonymousCredentials(), auth_request=auth_aiohttp.Request() + ) + try: + server.stalls[1] = 0.5 + session = AsyncResumableUploadSession( + upload_url=f"{base_url}/start", + config=ResumableUploadConfig(chunk_size=4), + transport=transport, + response_type=DummyResponse, + ) + result = await session.upload( + stream=b"01234567", timeout=0.2, retry=quick_retry + ) + assert isinstance(result, DummyResponse) + assert bytes(server.blobs[result.name]) == b"01234567" + assert server.commands == [ + "start", + "upload", + "query", + "upload", + "upload, finalize", + ] + + # Resume an upload the server already partially committed, then cancel. + server.blobs["seeded"] = bytearray(b"abcd") + server.commands.clear() + resumed = AsyncResumableUploadSession( + config=ResumableUploadConfig(chunk_size=4), + transport=transport, + response_type=DummyResponse, + ) + result = await resumed.resume( + upload_url=f"{base_url}/upload/seeded", stream=io.BytesIO(b"abcdefgh") + ) + await resumed.cancel() + assert isinstance(result, DummyResponse) + assert bytes(server.blobs["seeded"]) == b"abcdefgh" + assert server.commands == ["query", "upload, finalize", "cancel"] + assert server.cancelled == ["seeded"] + finally: + await transport.close() From d7b873b8b8ef7d30f3a081923cbf92cbceaafc46 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 1 Oct 2026 20:11:27 -0700 Subject: [PATCH 2/7] simplified implementation --- .../resumable_transfer/upload_async.py | 90 +++++++------ .../asyncio/test_resumable_transfer_async.py | 124 ++++-------------- 2 files changed, 77 insertions(+), 137 deletions(-) diff --git a/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py b/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py index 9dab3b749f9a..ce3f340b1c29 100644 --- a/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py +++ b/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py @@ -36,7 +36,6 @@ Mapping, Optional, Tuple, - Type, TypeVar, Union, cast, @@ -78,25 +77,11 @@ used by generated asynchronous REST clients. """ -# Timeouts and transport-level failures raised by the supported transports. -# ``AsyncAuthorizedSession`` surfaces aiohttp errors as google.auth exceptions. -_TIMEOUT_ERRORS: Tuple[Type[BaseException], ...] = (asyncio.TimeoutError,) -_TRANSPORT_ERRORS: Tuple[Type[BaseException], ...] = () -# Extra keyword arguments for every ``AsyncAuthorizedSession.request()`` call. -_AUTHORIZED_SESSION_REQUEST_KWARGS: Dict[str, Any] = {} -if _HAS_AIOHTTP: # pragma: no branch - _TIMEOUT_ERRORS += (aiohttp.ServerTimeoutError,) - _TRANSPORT_ERRORS += (aiohttp.ClientError,) -if _HAS_GOOGLE_AUTH_AIO: # pragma: no branch - _TIMEOUT_ERRORS += (auth_exceptions.TimeoutError,) - _TRANSPORT_ERRORS += (auth_exceptions.TransportError, auth_exceptions.ResponseError) - # google-auth >= 2.49.0 accepts ``total_attempts``; a single attempt disables - # its internal status-code retries so this module is the only retry layer. - _request_params = inspect.signature(AsyncAuthorizedSession.request).parameters - if "total_attempts" in _request_params: # pragma: no branch - _AUTHORIZED_SESSION_REQUEST_KWARGS["total_attempts"] = 1 -# Errors that may carry a chunk timeout out of the streaming retry loop. -_TIMEOUT_OR_RETRY_ERRORS = (exceptions.RetryError,) + _TIMEOUT_ERRORS +# google-auth >= 2.49.0 accepts ``total_attempts`` to cap AsyncAuthorizedSession's +# own status-code retries; older releases forward unknown kwargs to aiohttp. +_HAS_TOTAL_ATTEMPTS = _HAS_GOOGLE_AUTH_AIO and ( + "total_attempts" in inspect.signature(AsyncAuthorizedSession.request).parameters +) def _get_buffer_size(stream: object) -> Optional[int]: @@ -349,20 +334,33 @@ async def _send_request( Raises: exceptions.GoogleAPICallError: If the status is not 200 or 201. + asyncio.TimeoutError: If the attempt times out, whichever transport + is used. """ if _HAS_GOOGLE_AUTH_AIO and isinstance(transport, AsyncAuthorizedSession): - # ``request()`` is a coroutine returning a ``Response`` that must be - # closed explicitly. ``max_allowed_time`` keeps google-auth's - # wall-clock guard within the same per-attempt timeout. - response = await transport.request( - method, - url, - data=payload, - headers=headers, - timeout=timeout, - max_allowed_time=timeout, - **_AUTHORIZED_SESSION_REQUEST_KWARGS, + # google-auth applies its own overall deadline (``max_allowed_time``, + # 180s by default) on top of the per-request ``timeout``; use the + # caller's timeout for both. Where supported, a single attempt also + # disables google-auth's own status-code retries so that this module + # is the only retry layer. + retry_kwargs: Dict[str, Any] = ( + {"total_attempts": 1} if _HAS_TOTAL_ATTEMPTS else {} ) + try: + response = await transport.request( + method, + url, + data=payload, + headers=headers, + timeout=timeout, + max_allowed_time=timeout, + **retry_kwargs, + ) + except auth_exceptions.TimeoutError as exc: + # google-auth wraps timeouts in its own TimeoutError, which is + # unrelated to asyncio.TimeoutError; unwrap it so stall control + # sees the same timeout type as with aiohttp.ClientSession. + raise asyncio.TimeoutError(str(exc)) from exc try: status_code = response.status_code resp_headers: Mapping[str, str] = dict(response.headers) @@ -480,8 +478,8 @@ def _get_retry_predicate( 2. Protocol-recoverable errors during chunk transfer (``RECOVERABLE_STATUS_CODES`` and ``MissingStatusHeaderError``) and transport errors (``aiohttp.ClientError``, - ``asyncio.TimeoutError``, and their ``google.auth.exceptions`` - counterparts) always return ``True`` so the session + ``google.auth.exceptions.TransportError``, and + ``asyncio.TimeoutError``) always return ``True`` so the session can query server state and recover. Returns: @@ -499,7 +497,13 @@ def should_retry(exc: Exception) -> bool: ) ): return True - if isinstance(exc, _TIMEOUT_ERRORS + _TRANSPORT_ERRORS): + if isinstance(exc, asyncio.TimeoutError) or ( + _HAS_AIOHTTP and isinstance(exc, aiohttp.ClientError) + ): + return True + if _HAS_GOOGLE_AUTH_AIO and isinstance( + exc, (auth_exceptions.TransportError, auth_exceptions.ResponseError) + ): return True if ( custom_predicate is not None @@ -682,7 +686,13 @@ async def _transmit_chunk( t_elapsed = _monotonic_clock() - t_start except Exception as exc: self._enrich_exception(exc) - if isinstance(exc, _TIMEOUT_ERRORS): + if isinstance( + exc, + ( + asyncio.TimeoutError, + aiohttp.ServerTimeoutError, + ), + ): t_elapsed = _monotonic_clock() - t_start self._get_deadline_remaining() self._update_stall_control(0, t_elapsed) @@ -940,11 +950,17 @@ async def attempt_stream() -> AsyncGenerator[UploadProgress, None]: stream_gen = await retryable_stream() async for item in stream_gen: yield item - except _TIMEOUT_OR_RETRY_ERRORS as exc: + except ( + asyncio.TimeoutError, + aiohttp.ServerTimeoutError, + exceptions.RetryError, + ) as exc: timeout_exc = ( exc.__cause__ if isinstance(exc, exceptions.RetryError) else exc ) - if not isinstance(timeout_exc, _TIMEOUT_ERRORS): + if not isinstance( + timeout_exc, (asyncio.TimeoutError, aiohttp.ServerTimeoutError) + ): raise self._enrich_exception(timeout_exc) self._get_deadline_remaining() diff --git a/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py b/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py index 3a832b3d030c..318f4cab9af4 100644 --- a/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py +++ b/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py @@ -19,7 +19,6 @@ # mypy: disable-error-code="arg-type" import asyncio -import contextlib import datetime import inspect import io @@ -2433,8 +2432,8 @@ async def test_async_authorized_session_server_error_not_retried_internally() -> async def test_async_authorized_session_timeout_raises_transfer_stalled() -> None: """Verifies google.auth.exceptions.TimeoutError feeds stall control. - It is not a subclass of asyncio.TimeoutError, so it must be recognised - explicitly to surface as TransferStalledError. + It is not a subclass of asyncio.TimeoutError, so it is normalized to one + (keeping the original as the cause) to surface as TransferStalledError. """ transport, auth_request = _authorized_session( _auth_response(200, START_HEADERS), auth_exceptions.TimeoutError("timed out") @@ -2449,7 +2448,9 @@ async def test_async_authorized_session_timeout_raises_transfer_stalled() -> Non with pytest.raises(TransferStalledError) as exc_info: await session.upload(stream=b"data", retry=no_retry) - assert isinstance(exc_info.value.__cause__, auth_exceptions.TimeoutError) + cause = exc_info.value.__cause__ + assert isinstance(cause, asyncio.TimeoutError) + assert isinstance(cause.__cause__, auth_exceptions.TimeoutError) assert _adapter_commands(auth_request) == ["start", "upload, finalize"] @@ -2486,33 +2487,26 @@ async def test_async_authorized_session_transport_error_recovers() -> None: def test_async_retry_predicate_includes_google_auth_errors() -> None: - """Verifies google-auth transport/timeout errors are retryable, unlike its other errors.""" + """Verifies google-auth transport errors are retryable, unlike its other errors.""" session = AsyncResumableUploadSession(upload_url="https://api.example.com/start") for is_start in (True, False): predicate = session._get_retry_predicate(is_start=is_start) - assert predicate(auth_exceptions.TimeoutError()) is True assert predicate(auth_exceptions.TransportError()) is True assert predicate(auth_exceptions.ResponseError()) is True assert predicate(auth_exceptions.RefreshError()) is False # ===================================================================== -# 12. End-to-End Tests Against a Local HTTP Server +# 12. End-to-End Smoke Test Against a Local HTTP Server # ===================================================================== class ResumableUploadServer: - """Minimal in-memory resumable upload server. - - ``stalls`` maps a zero-based request ordinal to seconds to sleep before - answering that request with 408 without committing any bytes. - """ + """Minimal in-memory resumable upload server (start, upload, finalize).""" def __init__(self) -> None: self.blobs: Dict[str, bytearray] = {} self.commands: List[str] = [] - self.cancelled: List[str] = [] - self.stalls: Dict[int, float] = {} def make_app(self) -> "web.Application": app = web.Application() @@ -2522,12 +2516,7 @@ def make_app(self) -> "web.Application": async def handle(self, request: "web.Request") -> "web.Response": command = request.headers.get(common.HEADER_COMMAND, "") - payload = await request.read() - ordinal = len(self.commands) self.commands.append(command) - if ordinal in self.stalls: - await asyncio.sleep(self.stalls.pop(ordinal)) - return web.Response(status=408, text="stalled") if command == common._Command.START.value: blob = f"blob-{len(self.blobs)}" self.blobs[blob] = bytearray() @@ -2536,44 +2525,34 @@ async def handle(self, request: "web.Request") -> "web.Response": common.HEADER_URL: str(request.url.with_path(f"/upload/{blob}")), } return web.Response(headers=headers) - blob = request.match_info["blob"] - data = self.blobs[blob] - if command == common._Command.QUERY.value: - headers = { - common.HEADER_STATUS: common._Status.ACTIVE.value, - common.HEADER_SIZE_RECEIVED: str(len(data)), - } - return web.Response(headers=headers) - if command == common._Command.CANCEL.value: - self.cancelled.append(blob) - headers = {common.HEADER_STATUS: common._Status.CANCELLED.value} - return web.Response(headers=headers) + data = self.blobs[request.match_info["blob"]] if int(request.headers.get(common.HEADER_OFFSET, "-1")) != len(data): return web.Response(status=400, text="offset mismatch") - data.extend(payload) + data.extend(await request.read()) if common._Command.FINALIZE.value in command: headers = {common.HEADER_STATUS: common._Status.FINAL.value} - return web.json_response({"name": blob, "size": len(data)}, headers=headers) + body = {"name": request.match_info["blob"], "size": len(data)} + return web.json_response(body, headers=headers) return web.Response(headers={common.HEADER_STATUS: common._Status.ACTIVE.value}) -@contextlib.asynccontextmanager -async def serve_resumable_upload() -> AsyncIterator[Tuple[ResumableUploadServer, str]]: - """Serves a ``ResumableUploadServer`` locally, yielding it and its base URL.""" +@pytest.mark.asyncio +@pytest.mark.parametrize("authorized", [False, True], ids=["aiohttp", "google-auth"]) +async def test_async_upload_end_to_end(authorized: bool) -> None: + """Uploads over the wire with a real aiohttp.ClientSession / AsyncAuthorizedSession.""" server = ResumableUploadServer() async with TestServer(server.make_app()) as test_server: - yield server, str(test_server.make_url("")).rstrip("/") - - -@pytest.mark.asyncio -async def test_async_upload_end_to_end_aiohttp_client_session() -> None: - """Uploads over the wire with a real aiohttp.ClientSession.""" - async with serve_resumable_upload() as (server, base_url): async with aiohttp.ClientSession() as client: + transport: upload_async.AsyncTransport = client + if authorized: + transport = aio_sessions.AsyncAuthorizedSession( + aio_credentials.AnonymousCredentials(), + auth_request=auth_aiohttp.Request(client), + ) session = AsyncResumableUploadSession( - upload_url=f"{base_url}/start", + upload_url=str(test_server.make_url("/start")), config=ResumableUploadConfig(chunk_size=4), - transport=client, + transport=transport, response_type=DummyResponse, ) result = await session.upload(stream=b"0123456789") @@ -2581,58 +2560,3 @@ async def test_async_upload_end_to_end_aiohttp_client_session() -> None: assert isinstance(result, DummyResponse) assert bytes(server.blobs[result.name]) == b"0123456789" assert server.commands == ["start", "upload", "upload", "upload, finalize"] - - -@pytest.mark.asyncio -async def test_async_upload_end_to_end_authorized_session() -> None: - """Uploads, resumes, and cancels over the wire with the generated client transport. - - A real AsyncAuthorizedSession backed by google-auth's aiohttp adapter is - used. The first chunk request stalls past the per-attempt timeout, so - google-auth raises its own TimeoutError and api_core must query the - committed offset and re-send the chunk. - """ - quick_retry = google.api_core.retry.AsyncStreamingRetry(initial=0.01, maximum=0.01) - async with serve_resumable_upload() as (server, base_url): - transport = aio_sessions.AsyncAuthorizedSession( - aio_credentials.AnonymousCredentials(), auth_request=auth_aiohttp.Request() - ) - try: - server.stalls[1] = 0.5 - session = AsyncResumableUploadSession( - upload_url=f"{base_url}/start", - config=ResumableUploadConfig(chunk_size=4), - transport=transport, - response_type=DummyResponse, - ) - result = await session.upload( - stream=b"01234567", timeout=0.2, retry=quick_retry - ) - assert isinstance(result, DummyResponse) - assert bytes(server.blobs[result.name]) == b"01234567" - assert server.commands == [ - "start", - "upload", - "query", - "upload", - "upload, finalize", - ] - - # Resume an upload the server already partially committed, then cancel. - server.blobs["seeded"] = bytearray(b"abcd") - server.commands.clear() - resumed = AsyncResumableUploadSession( - config=ResumableUploadConfig(chunk_size=4), - transport=transport, - response_type=DummyResponse, - ) - result = await resumed.resume( - upload_url=f"{base_url}/upload/seeded", stream=io.BytesIO(b"abcdefgh") - ) - await resumed.cancel() - assert isinstance(result, DummyResponse) - assert bytes(server.blobs["seeded"]) == b"abcdefgh" - assert server.commands == ["query", "upload, finalize", "cancel"] - assert server.cancelled == ["seeded"] - finally: - await transport.close() From 9968e44a8e5b0813c993f36760f67d253cd18bed Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 1 Oct 2026 20:13:35 -0700 Subject: [PATCH 3/7] removed comment --- .../tests/asyncio/test_resumable_transfer_async.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py b/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py index 318f4cab9af4..3d5a33e4e45f 100644 --- a/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py +++ b/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py @@ -14,10 +14,6 @@ """Asynchronous tests for Resumable Upload protocol implementation.""" -# The transport doubles below only implement the parts of aiohttp.ClientSession -# that the session uses, so they do not satisfy the nominal AsyncTransport union. -# mypy: disable-error-code="arg-type" - import asyncio import datetime import inspect From 1c39dd28a67e629ce0fead7cd1275a8d2fea44a8 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 1 Oct 2026 20:17:21 -0700 Subject: [PATCH 4/7] removed redundant test --- .../asyncio/test_resumable_transfer_async.py | 71 +------------------ 1 file changed, 2 insertions(+), 69 deletions(-) diff --git a/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py b/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py index 3d5a33e4e45f..e697c55944c1 100644 --- a/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py +++ b/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py @@ -14,6 +14,8 @@ """Asynchronous tests for Resumable Upload protocol implementation.""" +# mypy: disable-error-code="arg-type" + import asyncio import datetime import inspect @@ -54,10 +56,7 @@ try: import aiohttp # noqa: F401 import google.auth.aio.transport # noqa: F401 - from aiohttp import web - from aiohttp.test_utils import TestServer from google.auth.aio import credentials as aio_credentials - from google.auth.aio.transport import aiohttp as auth_aiohttp from google.auth.aio.transport import sessions as aio_sessions GOOGLE_AUTH_AIO_INSTALLED = True @@ -2490,69 +2489,3 @@ def test_async_retry_predicate_includes_google_auth_errors() -> None: assert predicate(auth_exceptions.TransportError()) is True assert predicate(auth_exceptions.ResponseError()) is True assert predicate(auth_exceptions.RefreshError()) is False - - -# ===================================================================== -# 12. End-to-End Smoke Test Against a Local HTTP Server -# ===================================================================== - - -class ResumableUploadServer: - """Minimal in-memory resumable upload server (start, upload, finalize).""" - - def __init__(self) -> None: - self.blobs: Dict[str, bytearray] = {} - self.commands: List[str] = [] - - def make_app(self) -> "web.Application": - app = web.Application() - app.router.add_post("/start", self.handle) - app.router.add_post("/upload/{blob}", self.handle) - return app - - async def handle(self, request: "web.Request") -> "web.Response": - command = request.headers.get(common.HEADER_COMMAND, "") - self.commands.append(command) - if command == common._Command.START.value: - blob = f"blob-{len(self.blobs)}" - self.blobs[blob] = bytearray() - headers = { - common.HEADER_STATUS: common._Status.ACTIVE.value, - common.HEADER_URL: str(request.url.with_path(f"/upload/{blob}")), - } - return web.Response(headers=headers) - data = self.blobs[request.match_info["blob"]] - if int(request.headers.get(common.HEADER_OFFSET, "-1")) != len(data): - return web.Response(status=400, text="offset mismatch") - data.extend(await request.read()) - if common._Command.FINALIZE.value in command: - headers = {common.HEADER_STATUS: common._Status.FINAL.value} - body = {"name": request.match_info["blob"], "size": len(data)} - return web.json_response(body, headers=headers) - return web.Response(headers={common.HEADER_STATUS: common._Status.ACTIVE.value}) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("authorized", [False, True], ids=["aiohttp", "google-auth"]) -async def test_async_upload_end_to_end(authorized: bool) -> None: - """Uploads over the wire with a real aiohttp.ClientSession / AsyncAuthorizedSession.""" - server = ResumableUploadServer() - async with TestServer(server.make_app()) as test_server: - async with aiohttp.ClientSession() as client: - transport: upload_async.AsyncTransport = client - if authorized: - transport = aio_sessions.AsyncAuthorizedSession( - aio_credentials.AnonymousCredentials(), - auth_request=auth_aiohttp.Request(client), - ) - session = AsyncResumableUploadSession( - upload_url=str(test_server.make_url("/start")), - config=ResumableUploadConfig(chunk_size=4), - transport=transport, - response_type=DummyResponse, - ) - result = await session.upload(stream=b"0123456789") - - assert isinstance(result, DummyResponse) - assert bytes(server.blobs[result.name]) == b"0123456789" - assert server.commands == ["start", "upload", "upload", "upload, finalize"] From 2fd884880fe8315284b067d20b8d6f0406672b3e Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 2 Oct 2026 10:41:04 -0700 Subject: [PATCH 5/7] change message text --- .../google/api_core/resumable_transfer/upload_async.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py b/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py index ce3f340b1c29..350176b41982 100644 --- a/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py +++ b/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py @@ -357,9 +357,9 @@ async def _send_request( **retry_kwargs, ) except auth_exceptions.TimeoutError as exc: - # google-auth wraps timeouts in its own TimeoutError, which is - # unrelated to asyncio.TimeoutError; unwrap it so stall control - # sees the same timeout type as with aiohttp.ClientSession. + # google-auth raises its own TimeoutError (not an asyncio one, and + # not reliably chained to one); re-raise it as asyncio.TimeoutError + # so stall control sees the same type as with aiohttp.ClientSession. raise asyncio.TimeoutError(str(exc)) from exc try: status_code = response.status_code From 5bfe55881b02d2bcdff4ee067cb1ae8ec0fc26f9 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 2 Oct 2026 11:24:28 -0700 Subject: [PATCH 6/7] fix comment Co-authored-by: Anthonios Partheniou --- .../google/api_core/resumable_transfer/upload_async.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py b/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py index 350176b41982..ffb57ac15840 100644 --- a/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py +++ b/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py @@ -478,7 +478,8 @@ def _get_retry_predicate( 2. Protocol-recoverable errors during chunk transfer (``RECOVERABLE_STATUS_CODES`` and ``MissingStatusHeaderError``) and transport errors (``aiohttp.ClientError``, - ``google.auth.exceptions.TransportError``, and + ``google.auth.exceptions.TransportError``, + ``google.auth.exceptions.ResponseError``, and ``asyncio.TimeoutError``) always return ``True`` so the session can query server state and recover. From 69da04c17fb195c5e99697d97413b7267f70fe2a Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 2 Oct 2026 11:25:30 -0700 Subject: [PATCH 7/7] address PR comment --- .../api_core/resumable_transfer/upload_async.py | 12 ++++++------ .../tests/asyncio/test_resumable_transfer_async.py | 14 ++++++++++++-- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py b/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py index ffb57ac15840..d4c3e5da5ecb 100644 --- a/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py +++ b/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py @@ -356,17 +356,17 @@ async def _send_request( max_allowed_time=timeout, **retry_kwargs, ) + try: + status_code = response.status_code + resp_headers: Mapping[str, str] = dict(response.headers) + body = await response.read() + finally: + await response.close() except auth_exceptions.TimeoutError as exc: # google-auth raises its own TimeoutError (not an asyncio one, and # not reliably chained to one); re-raise it as asyncio.TimeoutError # so stall control sees the same type as with aiohttp.ClientSession. raise asyncio.TimeoutError(str(exc)) from exc - try: - status_code = response.status_code - resp_headers: Mapping[str, str] = dict(response.headers) - body = await response.read() - finally: - await response.close() else: # mypy cannot narrow the negative of the guarded check above. session = cast("aiohttp.ClientSession", transport) diff --git a/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py b/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py index e697c55944c1..a55941495d01 100644 --- a/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py +++ b/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py @@ -2424,14 +2424,22 @@ async def test_async_authorized_session_server_error_not_retried_internally() -> @pytest.mark.asyncio -async def test_async_authorized_session_timeout_raises_transfer_stalled() -> None: +@pytest.mark.parametrize("on_read", [False, True], ids=["request", "read"]) +async def test_async_authorized_session_timeout_raises_transfer_stalled( + on_read: bool, +) -> None: """Verifies google.auth.exceptions.TimeoutError feeds stall control. It is not a subclass of asyncio.TimeoutError, so it is normalized to one (keeping the original as the cause) to surface as TransferStalledError. """ + read_resp = _auth_response(200, {"X-Goog-Upload-Status": "final"}) + read_resp.read.side_effect = auth_exceptions.TimeoutError("read timed out") + chunk_outcome: Any = ( + read_resp if on_read else auth_exceptions.TimeoutError("timed out") + ) transport, auth_request = _authorized_session( - _auth_response(200, START_HEADERS), auth_exceptions.TimeoutError("timed out") + _auth_response(200, START_HEADERS), chunk_outcome ) session = AsyncResumableUploadSession( upload_url="https://api.example.com/start", transport=transport @@ -2447,6 +2455,8 @@ async def test_async_authorized_session_timeout_raises_transfer_stalled() -> Non assert isinstance(cause, asyncio.TimeoutError) assert isinstance(cause.__cause__, auth_exceptions.TimeoutError) assert _adapter_commands(auth_request) == ["start", "upload, finalize"] + if on_read: + read_resp.close.assert_awaited_once() @pytest.mark.asyncio