diff --git a/packages/google-api-core/google/api_core/exceptions.py b/packages/google-api-core/google/api_core/exceptions.py index df3e54e8f223..aa9898c5b9a6 100644 --- a/packages/google-api-core/google/api_core/exceptions.py +++ b/packages/google-api-core/google/api_core/exceptions.py @@ -446,6 +446,41 @@ class AsyncRestUnsupportedParameterError(NotImplementedError): pass +class ResumableTransferError(GoogleAPICallError): + """Base class for resumable transfer errors.""" + + upload_url: Optional[str] = None + chunk_size: Optional[int] = None + + def __init__( + self, + message: str, + *args, + upload_url: Optional[str] = None, + chunk_size: Optional[int] = None, + **kwargs, + ) -> None: + super().__init__(message, *args, **kwargs) + self.upload_url = upload_url + self.chunk_size = chunk_size + + +class TransferStalledError(ResumableTransferError): + """Raised when upload throughput stays below minimum rate past stall timeout.""" + + +class UnseekableStreamError(ResumableTransferError): + """Raised when server recovery requires rewinding a non-seekable stream.""" + + +class UploadCancelledError(ResumableTransferError): + """Raised when the upload is cancelled by the client or server.""" + + +class MissingStatusHeaderError(ResumableTransferError): + """Raised when server response lacks the required X-Goog-Upload-Status header.""" + + def exception_class_for_http_status(status_code): """Return the exception class for a specific HTTP status code. diff --git a/packages/google-api-core/google/api_core/resumable_transfer/__init__.py b/packages/google-api-core/google/api_core/resumable_transfer/__init__.py new file mode 100644 index 000000000000..9eabdcf65719 --- /dev/null +++ b/packages/google-api-core/google/api_core/resumable_transfer/__init__.py @@ -0,0 +1,49 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Resumable transfer library for Google APIs.""" + +from google.api_core.exceptions import ( + MissingStatusHeaderError, + ResumableTransferError, + TransferStalledError, + UnseekableStreamError, + UploadCancelledError, +) +from google.api_core.resumable_transfer.common import ( + DEFAULT_CHUNK_SIZE, + ProgressState, + ResumableUploadConfig, + UploadProgress, +) +from google.api_core.resumable_transfer.upload import ResumableUploadSession +from google.api_core.resumable_transfer.upload_async import ( + AsyncResumableUploadSession, + AsyncUploadOperation, +) + +__all__ = [ + "DEFAULT_CHUNK_SIZE", + "MissingStatusHeaderError", + "ProgressState", + "ResumableTransferError", + "TransferStalledError", + "UnseekableStreamError", + "UploadCancelledError", + "UploadProgress", + "ResumableUploadConfig", + "ResumableUploadSession", + "AsyncResumableUploadSession", + "AsyncUploadOperation", +] diff --git a/packages/google-api-core/google/api_core/resumable_transfer/common.py b/packages/google-api-core/google/api_core/resumable_transfer/common.py new file mode 100644 index 000000000000..e5d852814d2a --- /dev/null +++ b/packages/google-api-core/google/api_core/resumable_transfer/common.py @@ -0,0 +1,191 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Common constants, headers, and shared configuration for Resumable Upload protocol.""" + +import dataclasses +import datetime +import enum +from typing import Any, Mapping, Optional, Sequence, Tuple, Union + +import google.protobuf.message +import proto +from google.protobuf import json_format + +from google.api_core import exceptions + +# Default chunk size: 10 MiB +DEFAULT_CHUNK_SIZE = 10 * 1024 * 1024 + +# Default timeout in seconds for the initial start request +DEFAULT_START_TIMEOUT = 60.0 + +# Protocol Headers +HEADER_PROTOCOL = "X-Goog-Upload-Protocol" +HEADER_COMMAND = "X-Goog-Upload-Command" +HEADER_STATUS = "X-Goog-Upload-Status" +HEADER_URL = "X-Goog-Upload-URL" +HEADER_OFFSET = "X-Goog-Upload-Offset" +HEADER_SIZE_RECEIVED = "X-Goog-Upload-Size-Received" +HEADER_CONTENT_TYPE = "X-Goog-Upload-Header-Content-Type" +HEADER_CONTENT_LENGTH = "X-Goog-Upload-Header-Content-Length" +HEADER_CHUNK_GRANULARITY = "X-Goog-Upload-Chunk-Granularity" + +PROTOCOL_RESUMABLE = "resumable" + + +class _Command(str, enum.Enum): + """Protocol commands.""" + + START = "start" + UPLOAD = "upload" + FINALIZE = "finalize" + QUERY = "query" + CANCEL = "cancel" + + +class _Status(str, enum.Enum): + """Server upload status values.""" + + ACTIVE = "active" + FINAL = "final" + CANCELLED = "cancelled" + + +class ProgressState(str, enum.Enum): + """Progress notification state values.""" + + STARTED = "started" + UPLOADING = "uploading" + RECOVERING = "recovering" + OFFSET_RECEIVED = "offset received" + FINALIZED = "finalized" + + +@dataclasses.dataclass(frozen=True) +class UploadProgress: + """Upload progress notification payload. + + Attributes: + upload_url: The unique session URL for this upload. + chunk_size: The actual negotiated chunk size. + bytes_uploaded: The total confirmed bytes committed so far. + total_bytes: The total size of the stream in bytes, if known. + state: The current progress state. + """ + + upload_url: str + chunk_size: int + bytes_uploaded: int + total_bytes: Optional[int] + state: ProgressState + + +# HTTP status codes indicating transient retryable errors (Category 1) +RETRYABLE_STATUS_CODES = (408, 429, 500, 502, 503, 504) + +# HTTP status codes indicating state consistency errors requiring recovery (Category 2) +RECOVERABLE_STATUS_CODES = (400, 412, 416) + +# Exception types indicating unrecoverable terminal conditions (Category 3) +TERMINAL_ERRORS = ( + exceptions.DeadlineExceeded, + exceptions.TransferStalledError, + exceptions.UploadCancelledError, + exceptions.UnseekableStreamError, +) + + +@dataclasses.dataclass +class ResumableUploadConfig: + """Configuration options for a resumable upload. + + Attributes: + chunk_size: Size in bytes for each uploaded data chunk. Defaults to 10 MiB. + stall_minimum_rate: Minimum transfer rate in bytes per second. Defaults to 64 KiB/s. + stall_timeout: Stall duration threshold in seconds. Defaults to 120s. + headers: Additional HTTP headers dispatched exclusively with start request. + deadline: Optional overall wall-clock deadline for the entire upload process. + When set, each HTTP request timeout is trimmed to the remaining time before + the deadline, and DeadlineExceeded is raised immediately when the deadline + elapses (even on healthy streams). Timezone-aware datetimes (e.g., + ``datetime.now(timezone.utc) + timedelta(...)``) are recommended; + timezone-naive datetimes are assumed to be in local system time. When + None (default), transfer duration is governed by stall control + (stall_minimum_rate and stall_timeout), allowing healthy streams + transferring above the minimum rate to continue indefinitely. + """ + + chunk_size: int = DEFAULT_CHUNK_SIZE + stall_minimum_rate: int = 64 * 1024 + stall_timeout: float = 120.0 + headers: Optional[Union[Mapping[str, str], Sequence[Tuple[str, str]]]] = None + deadline: Optional[datetime.datetime] = None + + def __post_init__(self) -> None: + if self.deadline is not None: + self.deadline = self.deadline.astimezone(datetime.timezone.utc) + + @property + def start_headers(self) -> Optional[Sequence[Tuple[str, str]]]: + """Returns normalized additional headers for the start request.""" + if self.headers is None: + return None + if isinstance(self.headers, Mapping): + return list(self.headers.items()) + return list(self.headers) + + +def _format_response_payload( + response: Union[Any, bytes], + response_type: Optional[Any], +) -> Any: + """Formats raw response or bytes into protobuf or proto-plus message type if configured. + + Args: + response: Raw HTTP response object or response body bytes. + response_type: Deserializer callable, proto.Message class, or + google.protobuf.message.Message class or instance. + + Returns: + Deserialized protobuf message, or raw response bytes if + ``response_type`` is ``None``. + """ + content: bytes + if isinstance(response, bytes): + content = response + elif hasattr(response, "content"): + content = response.content + else: + content = bytes(response) + + if response_type is None: + return content + + from_json_fn = getattr(response_type, "from_json", None) + if callable(from_json_fn): + if isinstance(response_type, type) and issubclass(response_type, proto.Message): + return from_json_fn(content, ignore_unknown_fields=True) + return from_json_fn(content) + if isinstance(response_type, type) and issubclass( + response_type, google.protobuf.message.Message + ): + instance = response_type() + return json_format.Parse(content, instance, ignore_unknown_fields=True) + if isinstance(response_type, google.protobuf.message.Message): + return json_format.Parse(content, response_type, ignore_unknown_fields=True) + if callable(response_type): + return response_type(content) + + return content diff --git a/packages/google-api-core/google/api_core/resumable_transfer/upload.py b/packages/google-api-core/google/api_core/resumable_transfer/upload.py new file mode 100644 index 000000000000..5df0f6443721 --- /dev/null +++ b/packages/google-api-core/google/api_core/resumable_transfer/upload.py @@ -0,0 +1,1100 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Synchronous Resumable Upload session and helpers using requests.""" + +import datetime +import io +import logging +import time +from typing import ( + Any, + BinaryIO, + Callable, + Generator, + Iterable, + List, + Optional, + Tuple, + Union, +) + +import requests + +import google.api_core.retry +from google.api_core import exceptions +from google.api_core.resumable_transfer import common, upload_state +from google.api_core.resumable_transfer.common import ( + DEFAULT_START_TIMEOUT, + ResumableUploadConfig, + UploadProgress, + _format_response_payload, +) + +_LOGGER = logging.getLogger(__name__) +_monotonic_clock = time.monotonic + + +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) + if callable(getbuffer_fn): + return int(getbuffer_fn().nbytes) + return None + + +class _IterableReader(io.BytesIO): + """Wraps an Iterable[bytes] as a non-seekable binary stream. + + Ensure that chunks are pulled lazily from the underlying iterator + on each read() call rather than buffering the entire iterable into memory. + Inherits from io.BytesIO so static type checkers recognize instances as + BinaryIO natively. + """ + + def __init__(self, iterable: Iterable[bytes]) -> None: + super().__init__() + self._iterator = iter(iterable) + self._buffer = bytearray() + + def seekable(self) -> bool: + return False + + def tell(self) -> int: + raise OSError("Stream is not seekable") + + def read(self, size: Optional[int] = -1) -> bytes: + if size is None or size < 0: + for chunk in self._iterator: + self._buffer.extend(chunk) + result = bytes(self._buffer) + self._buffer.clear() + return result + + while len(self._buffer) < size: + try: + chunk = next(self._iterator) + self._buffer.extend(chunk) + except StopIteration: + break + result = bytes(self._buffer[:size]) + del self._buffer[:size] + return result + + +class ResumableUploadSession: + """Manages the full lifecycle of a resumable upload session.""" + + def __init__( + self, + upload_url: Optional[str] = None, + config: Optional[ResumableUploadConfig] = None, + transport: Optional[requests.Session] = None, + content_type: Optional[str] = None, + response_type: Optional[Any] = None, + start_retry: Optional[google.api_core.retry.Retry] = None, + start_timeout: float = DEFAULT_START_TIMEOUT, + ) -> None: + """Initializes a ResumableUploadSession. + + Args: + upload_url: The initial URL for the start request when starting a + new upload, or the pre-existing upload session URL when resuming. + config: Optional upload configuration parameters. Defaults to + ``ResumableUploadConfig()`` when ``None``. + transport: Optional requests session. 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. + response_type: Optional message class, callable deserializer, or + ``None``. When ``None``, raw response bytes are returned. + start_retry: Optional retry configuration (``google.api_core.retry.Retry``) + for the initial session creation request. When ``None``, the + default transient retry policy is used. Use this to customize + exponential backoff timing (such as ``Retry(initial=1.0, maximum=60.0)``) + or to supply a custom ``predicate`` function for API-specific transient + errors. A custom ``predicate`` replaces the default transient HTTP status + check (HTTP 408, 429, 500, 502, 503, and 504). Transport errors + (``ConnectionError``, ``ChunkedEncodingError``, and ``Timeout``) are + always retried. Terminal errors (``DeadlineExceeded``, + ``TransferStalledError``, ``UploadCancelledError``, and + ``UnseekableStreamError``) are never retried. + start_timeout: Timeout in seconds for the start request. Defaults to + ``60.0`` seconds. + """ + self._config = config or ResumableUploadConfig() + self._transport = transport + self._content_type = content_type + self._response_type = response_type + self._start_retry = start_retry + self._start_timeout = start_timeout + self._response: Optional[Any] = None + self._state = upload_state._ProtocolState( + upload_url=upload_url, + chunk_size=self._config.chunk_size, + ) + + # In-memory zero-copy buffer (never discard chunk until confirmed) + self._buffered_chunk: Optional[memoryview] = None + self._buffered_chunk_offset: int = 0 + self._buffered_chunk_is_last: bool = False + self._start_stream_offset: int = 0 + + # Stall control tracking via monotonic clock + self._aggregate_lag: float = 0.0 + self._stall_timeout_started: Optional[float] = None + self._needs_recovery: bool = False + + @property + def upload_url(self) -> Optional[str]: + """Optional[str]: The established upload session URL for this session. + + Once an upload is initiated, this is populated with the unique session + URL returned by the server in the ``x-goog-upload-url`` header (or updated + to the session URL passed to ``resume()``), and can be saved to resume an + interrupted upload later. + """ + return self._state.upload_url + + @property + def chunk_size(self) -> int: + """int: The negotiated chunk size.""" + return self._state.chunk_size + + @property + def response(self) -> Optional[Any]: + """Optional[Any]: The cached response message if finished.""" + return self._response + + @property + def bytes_uploaded(self) -> int: + """int: Confirmed number of bytes committed so far.""" + return self._state.bytes_uploaded + + @property + def finished(self) -> bool: + """bool: Whether the upload has completed successfully.""" + return self._state.finished + + def _reset_transfer_state(self) -> None: + """Resets per-transfer buffer, recovery, and stall control state.""" + self._response = None + self._buffered_chunk = None + self._buffered_chunk_offset = 0 + self._buffered_chunk_is_last = False + self._start_stream_offset = 0 + self._aggregate_lag = 0.0 + self._stall_timeout_started = None + self._needs_recovery = False + self._state._finished = False + self._state._invalid = False + + def _get_transport(self, transport: Optional[requests.Session]) -> requests.Session: + """Resolves the requests.Session transport. + + Args: + transport: Explicit requests session if provided. + + Returns: + The resolved requests session. + + Raises: + ValueError: If no requests session is available. + """ + sess = transport or self._transport + if sess is None: + raise ValueError("A requests.Session transport must be provided.") + return sess + + def _enrich_exception(self, exc: BaseException) -> None: + """Attaches session diagnostic metadata to an active exception. + + Args: + exc: Exception instance to augment with upload_url and chunk_size. + """ + setattr(exc, "upload_url", self.upload_url) + setattr(exc, "chunk_size", self.chunk_size) + + def _notify_progress( + self, + state: common.ProgressState, + progress_queue: Optional[List[UploadProgress]] = None, + ) -> None: + """Appends current upload status to the optional progress queue. + + Args: + state: ProgressState transition milestone. + progress_queue: Optional list buffering UploadProgress snapshots for generator consumers. + """ + if self.upload_url: + progress = UploadProgress( + upload_url=self.upload_url, + chunk_size=self.chunk_size, + bytes_uploaded=self._state.bytes_uploaded, + total_bytes=self._state.total_bytes, + state=state, + ) + if progress_queue is not None: + progress_queue.append(progress) + + def _get_deadline_remaining(self) -> Optional[float]: + """Calculates remaining seconds until the configured upload deadline. + + Returns: + Remaining seconds before deadline, or None if no deadline configured. + + Raises: + exceptions.DeadlineExceeded: If deadline has already elapsed. + """ + if self._config.deadline: + now = datetime.datetime.now(datetime.timezone.utc) + dl = self._config.deadline.astimezone(datetime.timezone.utc) + remaining = (dl - now).total_seconds() + if remaining <= 0: + raise exceptions.DeadlineExceeded( + f"Resumable upload deadline {self._config.deadline} exceeded." + ) + return remaining + return None + + def _get_start_timeout(self, timeout_override: Optional[float] = None) -> float: + """Computes timeout in seconds for start and control requests. + + Args: + timeout_override: Explicit timeout override in seconds. + + Returns: + Applicable timeout in seconds. + """ + remaining = self._get_deadline_remaining() + timeout = ( + timeout_override if timeout_override is not None else self._start_timeout + ) + if remaining is not None: + return min(timeout, remaining) + return timeout + + def _get_retry_predicate( + self, + is_start: bool = False, + custom_predicate: Optional[Callable[[Exception], bool]] = None, + ) -> Callable[[Exception], bool]: + """Returns a predicate function for determining if an exception is retryable. + + Args: + is_start: If True, only transient status codes (RETRYABLE_STATUS_CODES) + are retried. If False (transmitting and finalizing states), state + consistency errors (RECOVERABLE_STATUS_CODES and + MissingStatusHeaderError) are also retried via recovery. + custom_predicate: Optional callable taking an exception and returning True + if the error should be retried (from a user-supplied Retry instance). + Custom predicates replace the default transient HTTP status check + (``RETRYABLE_STATUS_CODES``: HTTP 408, 429, 500, 502, 503, and 504) + and are evaluated after protocol-enforced rules: + 1. Terminal errors (``TERMINAL_ERRORS``: ``DeadlineExceeded``, + ``TransferStalledError``, ``UploadCancelledError``, and + ``UnseekableStreamError``) always return ``False``. + 2. Protocol-recoverable errors during chunk transfer + (``RECOVERABLE_STATUS_CODES`` and ``MissingStatusHeaderError``) + and transport errors (``ConnectionError``, + ``ChunkedEncodingError``, ``Timeout``) always return ``True`` so + the session can query server state and recover. + + Returns: + A callable accepting an exception and returning a boolean. + """ + + def should_retry(exc: Exception) -> bool: + if isinstance(exc, common.TERMINAL_ERRORS): + return False + if not is_start and ( + isinstance(exc, exceptions.MissingStatusHeaderError) + or ( + isinstance(exc, exceptions.GoogleAPICallError) + and exc.code in common.RECOVERABLE_STATUS_CODES + ) + ): + return True + if isinstance(exc, requests.exceptions.RequestException): + if isinstance( + exc, + ( + requests.exceptions.ConnectionError, + requests.exceptions.ChunkedEncodingError, + requests.exceptions.Timeout, + ), + ): + return True + if ( + custom_predicate is not None + and custom_predicate is not google.api_core.retry.if_transient_error + ): + return bool(custom_predicate(exc)) + if isinstance(exc, exceptions.GoogleAPICallError): + return exc.code in common.RETRYABLE_STATUS_CODES + return False + + return should_retry + + def _get_retry( + self, retry_override: Optional[google.api_core.retry.Retry] = None + ) -> google.api_core.retry.Retry: + """Resolves unary Retry policy for start requests. + + Args: + retry_override: Optional unary Retry policy override for the start + request. Terminal errors (``DeadlineExceeded``, + ``TransferStalledError``, ``UploadCancelledError``, and + ``UnseekableStreamError``) are never retried. + + Returns: + Configured or default unary Retry instance. + """ + candidate = retry_override or self._start_retry + if candidate is not None: + return candidate.with_predicate( + self._get_retry_predicate( + is_start=True, custom_predicate=candidate._predicate + ) + ) + return google.api_core.retry.Retry( + predicate=self._get_retry_predicate(is_start=True) + ) + + def _get_streaming_retry( + self, + retry_override: Optional[google.api_core.retry.StreamingRetry] = None, + ) -> google.api_core.retry.StreamingRetry: + """Resolves the StreamingRetry policy for the chunk upload generator. + + Args: + retry_override: Optional StreamingRetry policy override for chunk + transmission. Protocol recovery is preserved automatically, and + terminal errors (``DeadlineExceeded``, ``TransferStalledError``, + ``UploadCancelledError``, and ``UnseekableStreamError``) are never + retried. + + Returns: + Configured or default StreamingRetry instance. + """ + if retry_override is not None: + wrapped_pred = self._get_retry_predicate( + is_start=False, custom_predicate=retry_override._predicate + ) + return retry_override.with_predicate(wrapped_pred) + return google.api_core.retry.StreamingRetry( + predicate=self._get_retry_predicate(is_start=False), + timeout=None, + ) + + def _compute_chunk_timeout( + self, data_len: int, timeout_override: Optional[float] = None + ) -> float: + """Computes the dynamic per-attempt chunk timeout based on stall control and deadlines. + + Args: + data_len: Length of the current chunk in bytes. + timeout_override: Optional per-attempt timeout ceiling in seconds. + + Returns: + Timeout in seconds for chunk transmission attempt. + """ + rate = self._config.stall_minimum_rate + expected_sec = data_len / rate if rate > 0 else 60.0 + + # Remaining time budget for this chunk across all retries before stalling. + # If prior chunks accumulated high aggregate_lag (or if expected_sec is tiny + # when stall_timeout is 0), expected_sec - _aggregate_lag + stall_timeout can + # drop to <= 0. Enforce a 1.0s minimum so the HTTP transport always receives + # a valid positive timeout instead of failing immediately with <= 0s. + min_chunk_timeout = 1.0 + next_chunk_timeout = max( + min_chunk_timeout, + expected_sec - self._aggregate_lag + self._config.stall_timeout, + ) + + if timeout_override is not None: + if rate > 0 and self._config.stall_timeout > 0: + per_attempt_timeout = min(timeout_override, next_chunk_timeout) + else: + per_attempt_timeout = timeout_override + else: + # Give a single attempt up to 2x expected_sec (bounded by next_chunk_timeout) + # so a hung socket fails fast enough to recover and retry before stalling. + # For a tiny final chunk (e.g. a few hundred bytes), 2x expected_sec is only + # a few milliseconds—shorter than an HTTP round-trip—so enforce a 5s minimum + # per-attempt timeout (e.g. 0.01s -> 5.0s, while 256.0s stays 256.0s). + min_attempt_timeout = 5.0 + per_attempt_timeout = max( + min_attempt_timeout, + min(next_chunk_timeout, 2.0 * expected_sec), + ) + + remaining = self._get_deadline_remaining() + if remaining is not None: + per_attempt_timeout = min(per_attempt_timeout, remaining) + + return per_attempt_timeout + + def _update_stall_control(self, data_len: int, t_elapsed: float) -> None: + """Updates aggregate transfer rate lag and enforces stall timeout and deadlines. + + Args: + data_len: Length of the transmitted chunk in bytes. + t_elapsed: Elapsed duration in seconds for chunk transmission. + + Raises: + exceptions.DeadlineExceeded: If upload deadline is exceeded. + exceptions.TransferStalledError: If transfer throughput stalls past configured timeout. + """ + if not (self._config.stall_minimum_rate and self._config.stall_timeout): + return + + rate = self._config.stall_minimum_rate + expected_sec = data_len / rate if rate > 0 else 0.0 + current_lag = t_elapsed - expected_sec + self._aggregate_lag = max(0.0, self._aggregate_lag + current_lag) + + if self._aggregate_lag > 0.0: + now = _monotonic_clock() + if self._stall_timeout_started is None: + self._stall_timeout_started = now - current_lag + if now - self._stall_timeout_started >= self._config.stall_timeout: + self._get_deadline_remaining() + raise exceptions.TransferStalledError( + f"Upload stalled: transfer rate remained below {rate} bytes/s " + f"for longer than {self._config.stall_timeout}s.", + upload_url=self.upload_url, + chunk_size=self.chunk_size, + ) + else: + self._stall_timeout_started = None + + def _reposition_stream_offset( + self, stream: Union[BinaryIO, Iterable[bytes]], received: int + ) -> int: + """Adjusts in-memory chunk buffer or seeks input stream to server offset. + + Args: + stream: The input data stream. + received: Confirmed byte offset committed on the server. + + Returns: + The confirmed server byte offset. + + Raises: + exceptions.UnseekableStreamError: If server offset precedes buffer and stream cannot be rewound. + """ + if self._buffered_chunk is not None: + chunk_start = self._buffered_chunk_offset + chunk_end = chunk_start + len(self._buffered_chunk) + if chunk_start <= received <= chunk_end: + discard_len = received - chunk_start + self._buffered_chunk = self._buffered_chunk[discard_len:] + self._buffered_chunk_offset = received + # When the server confirms receipt of the entire buffered chunk + # (received == chunk_end), slicing leaves a 0-length memoryview. + # Reset _buffered_chunk to None so the next upload attempt reads + # the next chunk from the stream instead of sending an empty buffer. + if len(self._buffered_chunk) == 0: + self._buffered_chunk = None + return received + + self._buffered_chunk = None + seekable_fn = getattr(stream, "seekable", None) + if callable(seekable_fn) and not seekable_fn(): + raise exceptions.UnseekableStreamError( + f"Stream is not seekable. Cannot recover upload to offset {received}.", + upload_url=self.upload_url, + chunk_size=self.chunk_size, + ) + try: + seek_fn = getattr(stream, "seek") + seek_fn(self._start_stream_offset + received) + except (OSError, AttributeError) as exc: + raise exceptions.UnseekableStreamError( + f"Failed to seek stream to offset {received}: {exc}", + upload_url=self.upload_url, + chunk_size=self.chunk_size, + ) from exc + + return received + + def _initiate( + self, + transport: requests.Session, + request_body: Union[str, bytes] = "", + size: Optional[int] = None, + progress_queue: Optional[List[UploadProgress]] = None, + content_type: Optional[str] = None, + ) -> str: + """Initiates the upload session by sending the start command. + + Args: + transport: The requests session. + request_body: JSON payload for initial start request. + size: Total size of payload in bytes, if known. + progress_queue: Optional list buffering UploadProgress snapshots. + content_type: Optional MIME type override of the payload. + + Returns: + The upload session URL. + """ + if content_type is not None: + self._content_type = content_type + + method, url, headers, payload = self._state.build_start_request( + body=request_body, + headers=self._config.start_headers, + content_type=self._content_type, + size=size, + ) + + def do_initiate() -> str: + req_timeout = self._get_start_timeout() + response = transport.request( + method, url, data=payload, headers=headers, timeout=req_timeout + ) + if not response.ok: + raise exceptions.from_http_response(response) + session_url = self._state.process_start_response( + response.status_code, response.headers + ) + return session_url + + retry_policy = self._get_retry() + retryable_initiate = retry_policy(do_initiate) + session_url = retryable_initiate() + self._notify_progress( + common.ProgressState.STARTED, progress_queue=progress_queue + ) + return session_url + + def _transmit_chunk( + self, + transport: requests.Session, + stream: Union[BinaryIO, Iterable[bytes]], + size: Optional[int], + progress_queue: Optional[List[UploadProgress]] = None, + timeout: Optional[float] = None, + ) -> requests.Response: + """Transmits a single data chunk attempt with stall control. + + Args: + transport: The requests session. + stream: The input data stream. + size: Total size of the stream in bytes, if known. + progress_queue: Optional list buffering UploadProgress snapshots. + timeout: Optional per-attempt timeout ceiling in seconds. + + Returns: + The HTTP response for the transmitted chunk. + """ + chunk_size = self._state.chunk_size + + # Retain active chunk in zero-copy buffer if not present. + # Ensure that EOF status (_buffered_chunk_is_last) is computed once + # when reading from the stream and preserved across _recover() retries. + # On partial server commit, _recover() slices _buffered_chunk in-place + # to the uncommitted tail. Preserving _buffered_chunk_is_last ensures + # that a sliced tail smaller than chunk_size is not prematurely + # treated as the final chunk when unread bytes remain in the stream. + if self._buffered_chunk is None: + read_fn = getattr(stream, "read") + raw_bytes = read_fn(chunk_size) + if not raw_bytes: + raw_bytes = b"" + self._buffered_chunk = memoryview(raw_bytes) + self._buffered_chunk_offset = self._state.bytes_uploaded + is_eof = len(raw_bytes) < chunk_size + if size is not None and self._state.bytes_uploaded + len(raw_bytes) >= size: + is_eof = True + self._buffered_chunk_is_last = is_eof + + data = self._buffered_chunk + data_len = len(data) + is_last = self._buffered_chunk_is_last + + method, url, headers, payload = self._state.build_chunk_request( + data=data, + is_last_chunk=is_last, + content_type=self._content_type, + ) + + per_attempt_timeout = self._compute_chunk_timeout( + data_len, timeout_override=timeout + ) + try: + t_start = _monotonic_clock() + resp = transport.request( + method, + url, + data=payload, + headers=headers, + timeout=per_attempt_timeout, + ) + t_elapsed = _monotonic_clock() - t_start + if not resp.ok: + raise exceptions.from_http_response(resp) + except requests.exceptions.Timeout as exc: + t_elapsed = _monotonic_clock() - t_start + self._enrich_exception(exc) + self._get_deadline_remaining() + self._update_stall_control(0, t_elapsed) + raise + except Exception as exc: + self._enrich_exception(exc) + raise + + self._update_stall_control(data_len, t_elapsed) + self._state.process_chunk_response(resp.status_code, resp.headers, data_len) + self._buffered_chunk = None + self._notify_progress( + common.ProgressState.FINALIZED + if self._state.finished + else common.ProgressState.UPLOADING, + progress_queue=progress_queue, + ) + return resp + + def _recover( + self, + transport: requests.Session, + stream: Union[BinaryIO, Iterable[bytes]], + progress_queue: Optional[List[UploadProgress]] = None, + ) -> requests.Response: + """Queries server for committed byte offset and adjusts buffer / stream. + + Args: + transport: The requests session. + stream: The input data stream. + progress_queue: Optional list buffering UploadProgress snapshots. + + Returns: + The HTTP response for the status query. + + Raises: + exceptions.UnseekableStreamError: If server offset precedes buffer and stream cannot be rewound. + exceptions.GoogleAPICallError: If query request fails on the server. + """ + method, url, headers, payload = self._state.build_query_request() + timeout = self._get_start_timeout() + resp = transport.request( + method, url, data=payload, headers=headers, timeout=timeout + ) + if not resp.ok: + raise exceptions.from_http_response(resp) + received = self._state.process_query_response(resp.status_code, resp.headers) + self._notify_progress( + common.ProgressState.OFFSET_RECEIVED, progress_queue=progress_queue + ) + self._reposition_stream_offset(stream, received) + return resp + + def cancel(self, transport: Optional[requests.Session] = None) -> None: + """Cancels the resumable upload session. + + Args: + transport: Optional requests session to use for dispatching cancellation. + + Raises: + ValueError: If no requests session is available. + GoogleAPICallError: If the cancellation request fails on the server. + """ + sess = self._get_transport(transport) + method, url, headers, payload = self._state.build_cancel_request() + timeout = self._get_start_timeout() + resp = sess.request(method, url, data=payload, headers=headers, timeout=timeout) + if not resp.ok: + raise exceptions.from_http_response(resp) + self._state.process_cancel_response(resp.status_code, resp.headers) + + def _transmit_all_chunks( + self, + transport: requests.Session, + stream_obj: Union[BinaryIO, Iterable[bytes]], + computed_size: Optional[int], + progress_queue: Optional[List[UploadProgress]] = None, + retry: Optional[google.api_core.retry.StreamingRetry] = None, + timeout: Optional[float] = None, + ) -> Generator[UploadProgress, None, None]: + """Transmits chunks until transfer completes, yielding buffered progress updates. + + Args: + transport: The requests session. + stream_obj: Binary stream yielding upload chunks. + computed_size: Total payload size in bytes if known. + progress_queue: Optional list buffering UploadProgress snapshots. + retry: Optional retry policy override for chunk transmission. Protocol + recovery is preserved automatically, and terminal errors + (``DeadlineExceeded``, ``TransferStalledError``, + ``UploadCancelledError``, and ``UnseekableStreamError``) are never + retried. + timeout: Optional per-attempt timeout ceiling in seconds. + + Yields: + UploadProgress snapshots for each transmission milestone. + + Raises: + ValueError: If upload concludes without a server response. + """ + if progress_queue is None: + progress_queue = [] + + while progress_queue: + yield progress_queue.pop(0) + + final_resp: Optional[requests.Response] = None + retry_policy = self._get_streaming_retry(retry_override=retry) + + def attempt_stream() -> Generator[UploadProgress, None, None]: + nonlocal final_resp + if self._needs_recovery: + _LOGGER.info( + "Recoverable error during chunk upload. Querying server offset." + ) + self._notify_progress( + common.ProgressState.RECOVERING, + progress_queue=progress_queue, + ) + recover_resp = self._recover( + transport, stream_obj, progress_queue=progress_queue + ) + if self._state.finished: + final_resp = recover_resp + self._needs_recovery = False + while progress_queue: + yield progress_queue.pop(0) + + while not self._state.finished and not self._state.invalid: + try: + final_resp = self._transmit_chunk( + transport, + stream_obj, + computed_size, + progress_queue=progress_queue, + timeout=timeout, + ) + except Exception as exc: + if retry_policy._predicate(exc): + self._needs_recovery = True + raise + while progress_queue: + yield progress_queue.pop(0) + + try: + retryable_stream = retry_policy(attempt_stream) + yield from retryable_stream() + except (requests.exceptions.Timeout, exceptions.RetryError) as exc: + timeout_exc = ( + exc.__cause__ if isinstance(exc, exceptions.RetryError) else exc + ) + if not isinstance(timeout_exc, requests.exceptions.Timeout): + raise + self._enrich_exception(timeout_exc) + self._get_deadline_remaining() + raise exceptions.TransferStalledError( + f"Upload stalled: chunk transfer timed out ({timeout_exc}).", + upload_url=self.upload_url, + chunk_size=self.chunk_size, + ) from timeout_exc + + if final_resp is None: + raise ValueError("Upload completed without receiving a final response.") + + self._response = _format_response_payload(final_resp, self._response_type) + + def upload( + self, + stream: Union[BinaryIO, bytes, Iterable[bytes]], + request_body: Union[str, bytes] = "", + size: Optional[int] = None, + transport: Optional[requests.Session] = None, + content_type: Optional[str] = None, + retry: Optional[google.api_core.retry.StreamingRetry] = None, + timeout: Optional[float] = None, + ) -> Any: + """Executes the resumable upload from start to completion. + + Args: + stream: Data payload to upload (file-like stream, bytes, or iterable of bytes). + request_body: Initial metadata payload sent with the start request. + size: Total stream size in bytes, if known. + transport: Optional requests session. + content_type: Optional MIME type of the stream payload. + retry: Optional retry configuration (``StreamingRetry``) for + chunk upload requests. Use this to customize exponential backoff + timing between chunk retries or to supply a custom ``predicate`` for + API-specific transient errors. A custom ``predicate`` replaces the + default transient HTTP status check (HTTP 408, 429, 500, 502, 503, + and 504). Transport errors and protocol recovery errors (HTTP 400, + 412, 416, and ``MissingStatusHeaderError``) always initiate server + offset recovery. Terminal errors (``DeadlineExceeded``, + ``TransferStalledError``, ``UploadCancelledError``, and + ``UnseekableStreamError``) are never retried. + timeout: Optional per-attempt timeout ceiling in seconds. + + Returns: + The final server response payload or deserialized response message. + + Raises: + ValueError: If transport is missing or upload completes without a response. + GoogleAPICallError: If an unrecoverable API error occurs. + """ + for _ in self.iter_upload( + stream=stream, + request_body=request_body, + size=size, + transport=transport, + content_type=content_type, + retry=retry, + timeout=timeout, + ): + pass + return self._response + + def iter_upload( + self, + stream: Union[BinaryIO, bytes, Iterable[bytes]], + request_body: Union[str, bytes] = "", + size: Optional[int] = None, + transport: Optional[requests.Session] = None, + content_type: Optional[str] = None, + retry: Optional[google.api_core.retry.StreamingRetry] = None, + timeout: Optional[float] = None, + ) -> Generator[UploadProgress, None, None]: + """Streams upload execution, yielding UploadProgress snapshots (PEP 255). + + Args: + stream: Data payload to upload (file-like stream, bytes, or iterable of bytes). + request_body: Initial metadata payload sent with the start request. + size: Total stream size in bytes, if known. + transport: Optional requests session. + content_type: Optional MIME type of the stream payload. + retry: Optional retry configuration (``StreamingRetry``) for + chunk upload requests. Use this to customize exponential backoff + timing between chunk retries or to supply a custom ``predicate`` for + API-specific transient errors. A custom ``predicate`` replaces the + default transient HTTP status check (HTTP 408, 429, 500, 502, 503, + and 504). Transport errors and protocol recovery errors (HTTP 400, + 412, 416, and ``MissingStatusHeaderError``) always initiate server + offset recovery. Terminal errors (``DeadlineExceeded``, + ``TransferStalledError``, ``UploadCancelledError``, and + ``UnseekableStreamError``) are never retried. + timeout: Optional per-attempt timeout ceiling in seconds. + + Yields: + UploadProgress snapshots for each chunk transmission milestone. + + Raises: + ValueError: If transport is missing or upload completes without a response. + GoogleAPICallError: If an unrecoverable API error occurs. + """ + sess = self._get_transport(transport) + if content_type is not None: + self._content_type = content_type + self._reset_transfer_state() + progress_queue: List[UploadProgress] = [] + try: + stream_obj, computed_size = self._prepare_stream(stream, size) + self._initiate( + transport=sess, + request_body=request_body, + size=computed_size, + progress_queue=progress_queue, + ) + yield from self._transmit_all_chunks( + sess, + stream_obj, + computed_size, + progress_queue=progress_queue, + retry=retry, + timeout=timeout, + ) + except Exception as exc: + self._enrich_exception(exc) + raise + + def resume( + self, + upload_url: str, + stream: Union[BinaryIO, bytes, Iterable[bytes]], + size: Optional[int] = None, + chunk_size: Optional[int] = None, + transport: Optional[requests.Session] = None, + retry: Optional[google.api_core.retry.StreamingRetry] = None, + timeout: Optional[float] = None, + ) -> Any: + """Resumes an existing upload from a saved upload URL. + + Args: + upload_url (str): The pre-existing upload session URL. + stream (Union[BinaryIO, bytes, Iterable[bytes]]): The data payload + to resume uploading from. + size (Optional[int]): Total size of the payload in bytes, if known. + chunk_size (Optional[int]): Optional chunk size override in bytes. + transport (Optional[requests.Session]): Optional requests session. + retry (Optional[google.api_core.retry.StreamingRetry]): Optional + retry configuration (``StreamingRetry``) for chunk upload + requests. Use this to customize exponential backoff timing + between chunk retries or to supply a custom ``predicate`` for + API-specific transient errors. A custom ``predicate`` replaces + the default transient HTTP status check (HTTP 408, 429, 500, + 502, 503, and 504). Transport errors and protocol recovery + errors (HTTP 400, 412, 416, and ``MissingStatusHeaderError``) + always initiate server offset recovery. Terminal errors + (``DeadlineExceeded``, ``TransferStalledError``, + ``UploadCancelledError``, and ``UnseekableStreamError``) are + never retried. + timeout (Optional[float]): Optional per-attempt timeout ceiling in + seconds. + + Returns: + The final server response payload or deserialized response message. + + Raises: + ValueError: If required arguments are missing or response not received. + GoogleAPICallError: If an unrecoverable API error occurs. + """ + for _ in self.iter_resume( + upload_url=upload_url, + stream=stream, + size=size, + chunk_size=chunk_size, + transport=transport, + retry=retry, + timeout=timeout, + ): + pass + return self._response + + def iter_resume( + self, + upload_url: str, + stream: Union[BinaryIO, bytes, Iterable[bytes]], + size: Optional[int] = None, + chunk_size: Optional[int] = None, + transport: Optional[requests.Session] = None, + retry: Optional[google.api_core.retry.StreamingRetry] = None, + timeout: Optional[float] = None, + ) -> Generator[UploadProgress, None, None]: + """Streams resumption of an upload, yielding UploadProgress snapshots. + + Args: + upload_url (str): The pre-existing upload session URL. + stream (Union[BinaryIO, bytes, Iterable[bytes]]): The data payload + to resume uploading from. + size (Optional[int]): Total size of the payload in bytes, if known. + chunk_size (Optional[int]): Optional chunk size override in bytes. + transport (Optional[requests.Session]): Optional requests session. + retry (Optional[google.api_core.retry.StreamingRetry]): Optional + retry configuration (``StreamingRetry``) for chunk upload + requests. Use this to customize exponential backoff timing + between chunk retries or to supply a custom ``predicate`` for + API-specific transient errors. A custom ``predicate`` replaces + the default transient HTTP status check (HTTP 408, 429, 500, + 502, 503, and 504). Transport errors and protocol recovery + errors (HTTP 400, 412, 416, and ``MissingStatusHeaderError``) + always initiate server offset recovery. Terminal errors + (``DeadlineExceeded``, ``TransferStalledError``, + ``UploadCancelledError``, and ``UnseekableStreamError``) are + never retried. + timeout (Optional[float]): Optional per-attempt timeout ceiling in + seconds. + + Yields: + UploadProgress snapshots for each chunk transmission milestone. + + Raises: + ValueError: If required arguments are missing or response not received. + GoogleAPICallError: If an unrecoverable API error occurs. + """ + 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.") + if stream is None: + raise ValueError("A data stream or payload must be provided to resume.") + + self._reset_transfer_state() + if chunk_size is not None: + self._state._chunk_size = chunk_size + + self._state._upload_url = actual_url + progress_queue: List[UploadProgress] = [] + try: + stream_obj, computed_size = self._prepare_stream(stream, size) + self._recover(sess, stream_obj, progress_queue=progress_queue) + yield from self._transmit_all_chunks( + sess, + stream_obj, + computed_size, + progress_queue=progress_queue, + retry=retry, + timeout=timeout, + ) + except Exception as exc: + self._enrich_exception(exc) + raise + + def _prepare_stream( + self, stream: Union[BinaryIO, bytes, Iterable[bytes]], size: Optional[int] + ) -> Tuple[Union[BinaryIO, Iterable[bytes]], Optional[int]]: + """Normalizes stream input into a readable stream object and determines stream length. + + Args: + stream: Input stream, bytes, or iterable of bytes. + size: Explicit total size in bytes, if known. + + Returns: + Tuple of (prepared stream object, computed total size). + """ + computed_size = size + if isinstance(stream, (str, dict)): + raise TypeError(f"Unsupported stream type: {type(stream)}") + if isinstance(stream, bytes): + stream_obj: Union[BinaryIO, Iterable[bytes]] = io.BytesIO(stream) + if computed_size is None: + computed_size = len(stream) + elif not hasattr(stream, "read") and isinstance(stream, Iterable): + stream_obj = _IterableReader(stream) + elif hasattr(stream, "read"): + stream_obj = stream + if computed_size is None: + computed_size = _get_buffer_size(stream_obj) + seekable_fn = getattr(stream_obj, "seekable", None) + tell_fn = getattr(stream_obj, "tell", None) + seek_fn = getattr(stream_obj, "seek", None) + if ( + computed_size is None + and callable(seekable_fn) + and seekable_fn() + and callable(tell_fn) + and callable(seek_fn) + ): + cur = tell_fn() + seek_fn(0, io.SEEK_END) + computed_size = tell_fn() - cur + seek_fn(cur) + else: + raise TypeError(f"Unsupported stream type: {type(stream)}") + + tell_fn = getattr(stream_obj, "tell", None) + if callable(tell_fn): + try: + self._start_stream_offset = tell_fn() + except (OSError, AttributeError): + self._start_stream_offset = 0 + + return stream_obj, computed_size 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 new file mode 100644 index 000000000000..d4c55a5f5f44 --- /dev/null +++ b/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py @@ -0,0 +1,1151 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Asynchronous Resumable Upload session and helpers using aiohttp.""" + +import asyncio +import datetime +import inspect +import io +import logging +import time +from typing import ( + Any, + AsyncGenerator, + AsyncIterable, + AsyncIterator, + Awaitable, + BinaryIO, + Callable, + Generator, + Generic, + Iterable, + List, + Mapping, + Optional, + Tuple, + TypeVar, + Union, +) + +try: + import aiohttp + + _HAS_AIOHTTP = True +except ImportError: # pragma: NO COVER + _HAS_AIOHTTP = False + +import google.api_core.retry +from google.api_core import exceptions +from google.api_core.resumable_transfer import common, upload_state +from google.api_core.resumable_transfer.common import ( + DEFAULT_START_TIMEOUT, + ResumableUploadConfig, + UploadProgress, + _format_response_payload, +) + +_LOGGER = logging.getLogger(__name__) +_monotonic_clock = time.monotonic + + +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) + if callable(getbuffer_fn): + return int(getbuffer_fn().nbytes) + return None + + +ResponseType = TypeVar("ResponseType") + + +class AsyncUploadOperation( + Generic[ResponseType], + Awaitable[ResponseType], + AsyncIterable[UploadProgress], +): + """Handle representing an asynchronous upload operation. + + Implements Awaitable[ResponseType] and AsyncIterable[UploadProgress] so + callers can either await the operation directly or iterate over progress + snapshots. + """ + + def __init__( + self, + progress_stream: AsyncIterator[UploadProgress], + session: "AsyncResumableUploadSession", + ) -> None: + """Initializes the upload operation handle. + + Args: + progress_stream: Async iterator yielding progress snapshots during the upload. + session: Underlying asynchronous resumable upload session. + """ + self._progress_stream = progress_stream + self._session = session + self._consumed = False + self._exception: Optional[BaseException] = None + + async def _consume(self) -> ResponseType: + """Consumes the upload stream to completion and returns the response.""" + if not self._consumed: + try: + async for _ in self._progress_stream: + pass + except BaseException as exc: + self._exception = exc + raise + finally: + self._consumed = True + if self._exception is not None: + raise self._exception + return self._session.response # type: ignore[return-value] + + def __await__(self) -> Generator[Any, None, ResponseType]: + """Awaits completion of the upload and returns the server response.""" + return self._consume().__await__() + + def __aiter__(self) -> AsyncIterator[UploadProgress]: + """Iterates over progress snapshots yielded during upload execution.""" + return self.progress() + + async def progress(self) -> AsyncIterator[UploadProgress]: + """Yields UploadProgress snapshots as the transfer progresses. + + Yields: + UploadProgress snapshots for each progress transition. + """ + if not self._consumed: + try: + async for item in self._progress_stream: + yield item + except BaseException as exc: + self._exception = exc + raise + finally: + self._consumed = True + + @property + def response(self) -> Optional[ResponseType]: + """The deserialized response message (or raw bytes), or None if in progress.""" + return self._session.response + + @property + def upload_url(self) -> Optional[str]: + """The session upload URL.""" + return self._session.upload_url + + @property + def chunk_size(self) -> int: + """The negotiated chunk size.""" + return self._session.chunk_size + + @property + def bytes_uploaded(self) -> int: + """Total confirmed bytes committed so far.""" + return self._session.bytes_uploaded + + +class AsyncResumableUploadSession: + """Manages the full lifecycle of an asynchronous resumable upload session.""" + + def __init__( + self, + upload_url: Optional[str] = None, + config: Optional[ResumableUploadConfig] = None, + transport: Optional[Any] = None, + content_type: Optional[str] = None, + response_type: Optional[Any] = None, + start_retry: Optional[google.api_core.retry.AsyncRetry] = None, + start_timeout: float = DEFAULT_START_TIMEOUT, + ) -> None: + """Initializes an AsyncResumableUploadSession. + + Args: + upload_url: The initial URL for the start request when starting a + 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 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. + response_type: Optional message class, callable deserializer, or + ``None``. When ``None``, raw response bytes are returned. + start_retry: Optional retry configuration (``google.api_core.retry.AsyncRetry``) + for the initial session creation request. When ``None``, the + default transient retry policy is used. Use this to customize + exponential backoff timing (such as ``AsyncRetry(initial=1.0, maximum=60.0)``) + or to supply a custom ``predicate`` function for API-specific transient + errors. A custom ``predicate`` replaces the default transient HTTP status + check (HTTP 408, 429, 500, 502, 503, and 504). Transport errors + (``aiohttp.ClientError``) are always retried. Terminal errors + (``DeadlineExceeded``, ``TransferStalledError``, + ``UploadCancelledError``, and ``UnseekableStreamError``) are never + retried. + start_timeout: Timeout in seconds for the start request. Defaults to + ``60.0`` seconds. + """ + self._config = config or ResumableUploadConfig() + self._transport = transport + self._content_type = content_type + self._response_type = response_type + self._start_retry = start_retry + self._start_timeout = start_timeout + self._response: Optional[Any] = None + self._state = upload_state._ProtocolState( + upload_url=upload_url, + chunk_size=self._config.chunk_size, + ) + + # In-memory zero-copy buffer + self._buffered_chunk: Optional[memoryview] = None + self._buffered_chunk_offset: int = 0 + self._buffered_chunk_is_last: bool = False + self._start_stream_offset: int = 0 + + # Stall control tracking via monotonic clock + self._aggregate_lag: float = 0.0 + self._stall_timeout_started: Optional[float] = None + self._needs_recovery: bool = False + + @property + def upload_url(self) -> Optional[str]: + """Optional[str]: The established upload session URL for this session. + + Once an upload is initiated, this is populated with the unique session + URL returned by the server in the ``x-goog-upload-url`` header (or updated + to the session URL passed to ``resume()``), and can be saved to resume an + interrupted upload later. + """ + return self._state.upload_url + + @property + def chunk_size(self) -> int: + """int: The negotiated chunk size.""" + return self._state.chunk_size + + @property + def response(self) -> Optional[Any]: + """Optional[Any]: The cached response message if finished.""" + return self._response + + @property + def bytes_uploaded(self) -> int: + """int: Confirmed number of bytes committed so far.""" + return self._state.bytes_uploaded + + @property + def finished(self) -> bool: + """bool: Whether the upload has completed successfully.""" + return self._state.finished + + def _reset_transfer_state(self) -> None: + """Resets per-transfer buffer, recovery, and stall control state.""" + self._response = None + self._buffered_chunk = None + self._buffered_chunk_offset = 0 + self._buffered_chunk_is_last = False + self._start_stream_offset = 0 + self._aggregate_lag = 0.0 + self._stall_timeout_started = None + self._needs_recovery = False + self._state._finished = False + self._state._invalid = False + + def _ensure_aiohttp(self) -> None: + """Validates that aiohttp is installed and accessible. + + Raises: + ImportError: If aiohttp is not installed. + """ + if not _HAS_AIOHTTP: + raise ImportError( + "The aiohttp library is required to use AsyncResumableUploadSession. " + "Please install google-api-core[async_rest]." + ) + + def _enrich_exception(self, exc: BaseException) -> None: + """Attaches session diagnostic metadata to an active exception. + + Args: + exc: Exception instance to augment with upload_url and chunk_size. + """ + setattr(exc, "upload_url", self.upload_url) + setattr(exc, "chunk_size", self.chunk_size) + + def _notify_progress( + self, + state: common.ProgressState, + queue: Optional[List[UploadProgress]] = None, + ) -> None: + """Appends current upload status to the optional progress queue. + + Args: + state: ProgressState transition milestone. + queue: Optional list to receive progress event. + """ + if self.upload_url: + progress = UploadProgress( + upload_url=self.upload_url, + chunk_size=self.chunk_size, + bytes_uploaded=self._state.bytes_uploaded, + total_bytes=self._state.total_bytes, + state=state, + ) + if queue is not None: + queue.append(progress) + + def _get_deadline_remaining(self) -> Optional[float]: + """Calculates remaining seconds until the configured upload deadline. + + Returns: + Remaining seconds before deadline, or None if no deadline configured. + + Raises: + exceptions.DeadlineExceeded: If deadline has already elapsed. + """ + if self._config.deadline: + now = datetime.datetime.now(datetime.timezone.utc) + dl = self._config.deadline.astimezone(datetime.timezone.utc) + remaining = (dl - now).total_seconds() + if remaining <= 0: + raise exceptions.DeadlineExceeded( + f"Resumable upload deadline {self._config.deadline} exceeded." + ) + return remaining + return None + + def _get_start_timeout(self, timeout_override: Optional[float] = None) -> float: + """Computes timeout in seconds for start and control requests. + + Args: + timeout_override: Explicit timeout override in seconds. + + Returns: + Applicable timeout in seconds. + """ + remaining = self._get_deadline_remaining() + timeout = ( + timeout_override if timeout_override is not None else self._start_timeout + ) + if remaining is not None: + return min(timeout, remaining) + return timeout + + def _get_retry_predicate( + self, + is_start: bool = False, + custom_predicate: Optional[Callable[[Exception], bool]] = None, + ) -> Callable[[Exception], bool]: + """Returns a predicate function for determining if an exception is retryable. + + Args: + is_start: If True, only transient status codes (RETRYABLE_STATUS_CODES) + are retried. If False (transmitting and finalizing states), state + consistency errors (RECOVERABLE_STATUS_CODES and + MissingStatusHeaderError) are also retried via recovery. + custom_predicate: Optional callable taking an exception and returning True + if the error should be retried (from a user-supplied AsyncRetry instance). + Custom predicates replace the default transient HTTP status check + (``RETRYABLE_STATUS_CODES``: HTTP 408, 429, 500, 502, 503, and 504) + and are evaluated after protocol-enforced rules: + 1. Terminal errors (``TERMINAL_ERRORS``: ``DeadlineExceeded``, + ``TransferStalledError``, ``UploadCancelledError``, and + ``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 + can query server state and recover. + + Returns: + A callable accepting an exception and returning a boolean. + """ + + def should_retry(exc: Exception) -> bool: + if isinstance(exc, common.TERMINAL_ERRORS): + return False + if not is_start and ( + isinstance(exc, exceptions.MissingStatusHeaderError) + or ( + isinstance(exc, exceptions.GoogleAPICallError) + and exc.code in common.RECOVERABLE_STATUS_CODES + ) + ): + return True + if isinstance(exc, asyncio.TimeoutError) or ( + _HAS_AIOHTTP and isinstance(exc, aiohttp.ClientError) + ): + return True + if ( + custom_predicate is not None + and custom_predicate is not google.api_core.retry.if_transient_error + ): + return bool(custom_predicate(exc)) + if isinstance(exc, exceptions.GoogleAPICallError): + return exc.code in common.RETRYABLE_STATUS_CODES + return False + + return should_retry + + def _get_async_retry( + self, retry_override: Optional[google.api_core.retry.AsyncRetry] = None + ) -> google.api_core.retry.AsyncRetry: + """Resolves unary AsyncRetry policy for start requests. + + Args: + retry_override: Optional unary AsyncRetry policy override for the start + request. Terminal errors (``DeadlineExceeded``, + ``TransferStalledError``, ``UploadCancelledError``, and + ``UnseekableStreamError``) are never retried. + + Returns: + Configured or default unary AsyncRetry instance. + """ + candidate = retry_override or self._start_retry + if candidate is not None: + return candidate.with_predicate( + self._get_retry_predicate( + is_start=True, custom_predicate=candidate._predicate + ) + ) + return google.api_core.retry.AsyncRetry( + predicate=self._get_retry_predicate(is_start=True) + ) + + def _get_async_streaming_retry( + self, + retry_override: Optional[google.api_core.retry.AsyncStreamingRetry] = None, + ) -> google.api_core.retry.AsyncStreamingRetry: + """Resolves the AsyncStreamingRetry policy for the chunk upload generator. + + Args: + retry_override: Optional AsyncStreamingRetry policy override for chunk + transmission. Protocol recovery is preserved automatically, and + terminal errors (``DeadlineExceeded``, ``TransferStalledError``, + ``UploadCancelledError``, and ``UnseekableStreamError``) are never + retried. + + Returns: + Configured or default AsyncStreamingRetry instance. + """ + if retry_override is not None: + wrapped_pred = self._get_retry_predicate( + is_start=False, custom_predicate=retry_override._predicate + ) + return retry_override.with_predicate(wrapped_pred) + return google.api_core.retry.AsyncStreamingRetry( + predicate=self._get_retry_predicate(is_start=False), + timeout=None, + ) + + async def _initiate( + self, + transport: Any, + request_body: Union[str, bytes] = "", + size: Optional[int] = None, + progress_queue: Optional[List[UploadProgress]] = None, + content_type: Optional[str] = None, + ) -> str: + """Initiates the upload session asynchronously. + + Args: + transport: The aiohttp client session. + 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. + content_type: Optional MIME type override of the payload. + + Returns: + The negotiated upload session URL. + """ + self._ensure_aiohttp() + if content_type is not None: + self._content_type = content_type + + method, url, headers, payload = self._state.build_start_request( + body=request_body, + headers=self._config.start_headers, + content_type=self._content_type, + size=size, + ) + + 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 + + retry_policy = self._get_async_retry() + retryable_initiate = retry_policy(do_initiate) + session_url = await retryable_initiate() + self._notify_progress(common.ProgressState.STARTED, progress_queue) + return session_url + + async def _transmit_chunk( + self, + transport: Any, + reader_fn: Callable[[int], Awaitable[bytes]], + size: Optional[int], + progress_queue: Optional[List[UploadProgress]] = None, + timeout: Optional[float] = None, + ) -> Tuple[int, Mapping[str, str], bytes]: + """Transmits the next data chunk asynchronously with stall control. + + Args: + transport: The aiohttp client session. + reader_fn: Async callable returning chunk bytes. + size: Total stream size in bytes, if known. + progress_queue: Optional queue to receive progress updates. + timeout: Optional per-attempt timeout ceiling in seconds. + + Returns: + Tuple of (status code, headers mapping, response body bytes). + + Raises: + TransferStalledError: If chunk transfer throughput stalls. + DeadlineExceeded: If upload deadline is reached. + GoogleAPICallError: If chunk upload encounters an error. + """ + chunk_size = self._state.chunk_size + + # Retain active chunk in zero-copy buffer if not present. + # Ensure that EOF status (_buffered_chunk_is_last) is computed once + # when reading from the stream and preserved across _recover() retries. + # On partial server commit, _recover() slices _buffered_chunk in-place + # to the uncommitted tail. Preserving _buffered_chunk_is_last ensures + # that a sliced tail smaller than chunk_size is not prematurely + # treated as the final chunk when unread bytes remain in the stream. + if self._buffered_chunk is None: + raw_bytes = await reader_fn(chunk_size) + if not raw_bytes: + raw_bytes = b"" + self._buffered_chunk = memoryview(raw_bytes) + self._buffered_chunk_offset = self._state.bytes_uploaded + is_eof = len(raw_bytes) < chunk_size + if size is not None and self._state.bytes_uploaded + len(raw_bytes) >= size: + is_eof = True + self._buffered_chunk_is_last = is_eof + + data = self._buffered_chunk + data_len = len(data) + is_last = self._buffered_chunk_is_last + + method, url, headers, payload = self._state.build_chunk_request( + data=data, + is_last_chunk=is_last, + content_type=self._content_type, + ) + + 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( + 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 + t_elapsed = _monotonic_clock() - t_start + except Exception as exc: + self._enrich_exception(exc) + if isinstance( + exc, + ( + asyncio.TimeoutError, + aiohttp.ServerTimeoutError, + ), + ): + t_elapsed = _monotonic_clock() - t_start + self._get_deadline_remaining() + self._update_stall_control(0, t_elapsed) + raise + + self._update_stall_control(data_len, t_elapsed) + + self._state.process_chunk_response(status_code, resp_headers, data_len) + self._buffered_chunk = None + self._notify_progress( + common.ProgressState.FINALIZED + if self._state.finished + else common.ProgressState.UPLOADING, + progress_queue, + ) + return status_code, resp_headers, resp_body + + def _compute_chunk_timeout( + self, data_len: int, timeout_override: Optional[float] = None + ) -> float: + """Computes the dynamic per-attempt chunk timeout based on stall control and deadlines. + + Args: + data_len: Length of the current chunk in bytes. + timeout_override: Optional per-attempt timeout ceiling in seconds. + + Returns: + Timeout in seconds for chunk transmission attempt. + """ + rate = self._config.stall_minimum_rate + expected_sec = data_len / rate if rate > 0 else 60.0 + + # Remaining time budget for this chunk across all retries before stalling. + # If prior chunks accumulated high aggregate_lag (or if expected_sec is tiny + # when stall_timeout is 0), expected_sec - _aggregate_lag + stall_timeout can + # drop to <= 0. Enforce a 1.0s minimum so the HTTP transport always receives + # a valid positive timeout instead of failing immediately with <= 0s. + min_chunk_timeout = 1.0 + next_chunk_timeout = max( + min_chunk_timeout, + expected_sec - self._aggregate_lag + self._config.stall_timeout, + ) + + if timeout_override is not None: + if rate > 0 and self._config.stall_timeout > 0: + per_attempt_timeout = min(timeout_override, next_chunk_timeout) + else: + per_attempt_timeout = timeout_override + else: + # Give a single attempt up to 2x expected_sec (bounded by next_chunk_timeout) + # so a hung socket fails fast enough to recover and retry before stalling. + # For a tiny final chunk (e.g. a few hundred bytes), 2x expected_sec is only + # a few milliseconds—shorter than an HTTP round-trip—so enforce a 5s minimum + # per-attempt timeout (e.g. 0.01s -> 5.0s, while 256.0s stays 256.0s). + min_attempt_timeout = 5.0 + per_attempt_timeout = max( + min_attempt_timeout, + min(next_chunk_timeout, 2.0 * expected_sec), + ) + + remaining = self._get_deadline_remaining() + if remaining is not None: + per_attempt_timeout = min(per_attempt_timeout, remaining) + + return per_attempt_timeout + + def _update_stall_control(self, data_len: int, t_elapsed: float) -> None: + """Updates aggregate transfer rate lag and enforces stall timeout and deadlines. + + Args: + data_len: Length of the transmitted chunk in bytes. + t_elapsed: Elapsed duration in seconds for chunk transmission. + + Raises: + exceptions.DeadlineExceeded: If upload deadline is exceeded. + exceptions.TransferStalledError: If transfer throughput stalls past configured timeout. + """ + if not (self._config.stall_minimum_rate and self._config.stall_timeout): + return + + rate = self._config.stall_minimum_rate + expected_sec = data_len / rate if rate > 0 else 0.0 + current_lag = t_elapsed - expected_sec + self._aggregate_lag = max(0.0, self._aggregate_lag + current_lag) + if self._aggregate_lag > 0.0: + now = _monotonic_clock() + if self._stall_timeout_started is None: + self._stall_timeout_started = now - current_lag + if now - self._stall_timeout_started >= self._config.stall_timeout: + self._get_deadline_remaining() + raise exceptions.TransferStalledError( + f"Upload stalled: transfer rate remained below {rate} bytes/s for longer than {self._config.stall_timeout}s.", + upload_url=self.upload_url, + chunk_size=self.chunk_size, + ) + else: + self._stall_timeout_started = None + + async def _recover( + self, + transport: Any, + 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. + stream_obj: Underlying stream object to rewind if seekable. + progress_queue: Optional queue to receive progress updates. + + Returns: + Tuple of (status code, headers mapping, response body bytes). + + Raises: + exceptions.UnseekableStreamError: If server offset precedes buffer and stream cannot be rewound. + 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 + + received = self._state.process_query_response(status_code, resp_headers) + self._notify_progress(common.ProgressState.OFFSET_RECEIVED, progress_queue) + + if self._buffered_chunk is not None: + chunk_start = self._buffered_chunk_offset + chunk_end = chunk_start + len(self._buffered_chunk) + if chunk_start <= received <= chunk_end: + discard_len = received - chunk_start + self._buffered_chunk = self._buffered_chunk[discard_len:] + self._buffered_chunk_offset = received + # When the server confirms receipt of the entire buffered chunk + # (received == chunk_end), slicing leaves a 0-length memoryview. + # Reset _buffered_chunk to None so the next upload attempt reads + # the next chunk from the stream instead of sending an empty buffer. + if len(self._buffered_chunk) == 0: + self._buffered_chunk = None + return status_code, resp_headers, body + + self._buffered_chunk = None + seek_fn = getattr(stream_obj, "seek", None) + if callable(seek_fn): + seekable_fn = getattr(stream_obj, "seekable", None) + if callable(seekable_fn) and not seekable_fn(): + raise exceptions.UnseekableStreamError( + f"Stream is not seekable. Cannot recover upload to offset {received}.", + upload_url=self.upload_url, + chunk_size=self.chunk_size, + ) + try: + seek_fn(self._start_stream_offset + received) + return status_code, resp_headers, body + except (OSError, AttributeError) as exc: + raise exceptions.UnseekableStreamError( + f"Failed to seek stream to offset {received}: {exc}", + upload_url=self.upload_url, + chunk_size=self.chunk_size, + ) from exc + + raise exceptions.UnseekableStreamError( + f"Server offset {received} precedes active buffer. Stream cannot be rewound.", + upload_url=self.upload_url, + chunk_size=self.chunk_size, + ) + + async def cancel(self, transport: Optional[Any] = None) -> None: + """Cancels the resumable upload session asynchronously. + + Args: + transport: Optional aiohttp client session. + + 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.") + 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) + + async def _transmit_all_chunks( + self, + transport: Any, + reader_fn: Callable[[int], Awaitable[bytes]], + computed_size: Optional[int], + progress_queue: Optional[List[UploadProgress]] = None, + stream_obj: Optional[object] = None, + retry: Optional[google.api_core.retry.AsyncStreamingRetry] = None, + timeout: Optional[float] = None, + ) -> AsyncGenerator[UploadProgress, None]: + """Transmits chunks until completion using a single outer AsyncStreamingRetry coordinator. + + Args: + transport: The aiohttp client session. + reader_fn: Async callable returning chunk bytes. + computed_size: Total stream size in bytes, if known. + progress_queue: Optional list receiving UploadProgress snapshots. + stream_obj: Underlying stream object for recovery seeking. + retry: Optional retry policy override for chunk transmission. Protocol + recovery is preserved automatically, and terminal errors + (``DeadlineExceeded``, ``TransferStalledError``, + ``UploadCancelledError``, and ``UnseekableStreamError``) are never + retried. + timeout: Optional per-attempt timeout ceiling in seconds. + + Yields: + UploadProgress snapshots for each chunk transmission milestone. + """ + final_resp_tuple: Optional[Tuple[int, Mapping[str, str], bytes]] = None + retry_policy = self._get_async_streaming_retry(retry_override=retry) + + async def attempt_stream() -> AsyncGenerator[UploadProgress, None]: + nonlocal final_resp_tuple + if self._needs_recovery: + _LOGGER.info( + "Recoverable error during async chunk upload. Querying server offset." + ) + self._notify_progress(common.ProgressState.RECOVERING, progress_queue) + recover_tuple = await self._recover( + transport, stream_obj, progress_queue=progress_queue + ) + if self._state.finished: + final_resp_tuple = recover_tuple + self._needs_recovery = False + while progress_queue: + yield progress_queue.pop(0) + + while not self._state.finished and not self._state.invalid: + try: + final_resp_tuple = await self._transmit_chunk( + transport, + reader_fn, + computed_size, + progress_queue, + timeout=timeout, + ) + except Exception as exc: + if retry_policy._predicate(exc): + self._needs_recovery = True + raise + while progress_queue: + yield progress_queue.pop(0) + + try: + retryable_stream = retry_policy(attempt_stream) + stream_gen = await retryable_stream() + async for item in stream_gen: + yield item + 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, (asyncio.TimeoutError, aiohttp.ServerTimeoutError) + ): + raise + self._enrich_exception(timeout_exc) + self._get_deadline_remaining() + raise exceptions.TransferStalledError( + f"Upload stalled: chunk transfer timed out ({timeout_exc}).", + upload_url=self.upload_url, + chunk_size=self.chunk_size, + ) from timeout_exc + + if final_resp_tuple is None: + raise ValueError("Upload completed without receiving a final response.") + + _, _, body_bytes = final_resp_tuple + self._response = _format_response_payload(body_bytes, self._response_type) + + def upload( + self, + stream: Union[AsyncIterable[bytes], BinaryIO, bytes, Iterable[bytes]], + request_body: Union[str, bytes] = "", + size: Optional[int] = None, + transport: Optional[Any] = None, + content_type: Optional[str] = None, + retry: Optional[google.api_core.retry.AsyncStreamingRetry] = None, + timeout: Optional[float] = None, + ) -> AsyncUploadOperation: + """Initiates and executes upload asynchronously, returning an AsyncUploadOperation. + + Args: + 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. + 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 + supply a custom ``predicate`` for API-specific transient errors. + A custom ``predicate`` replaces the default transient HTTP status + check (HTTP 408, 429, 500, 502, 503, and 504). Transport errors and + protocol recovery errors (HTTP 400, 412, 416, and + ``MissingStatusHeaderError``) always initiate server offset + recovery. Terminal errors (``DeadlineExceeded``, + ``TransferStalledError``, ``UploadCancelledError``, and + ``UnseekableStreamError``) are never retried. + timeout: Optional per-attempt timeout ceiling in seconds. + + Returns: + An AsyncUploadOperation handle representing the transfer. + + Raises: + 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.") + + if content_type is not None: + self._content_type = content_type + + self._reset_transfer_state() + progress_queue: List[UploadProgress] = [] + reader_fn, computed_size, stream_obj = self._prepare_async_reader(stream, size) + + async def _run() -> AsyncGenerator[UploadProgress, None]: + try: + await self._initiate( + transport=sess, + request_body=request_body, + size=computed_size, + progress_queue=progress_queue, + ) + while progress_queue: + yield progress_queue.pop(0) + + async for item in self._transmit_all_chunks( + sess, + reader_fn, + computed_size, + progress_queue=progress_queue, + stream_obj=stream_obj, + retry=retry, + timeout=timeout, + ): + yield item + except BaseException as exc: + self._enrich_exception(exc) + raise + + return AsyncUploadOperation(progress_stream=_run(), session=self) + + def resume( + self, + upload_url: str, + stream: Union[AsyncIterable[bytes], BinaryIO, bytes, Iterable[bytes]], + size: Optional[int] = None, + chunk_size: Optional[int] = None, + transport: Optional[Any] = None, + retry: Optional[google.api_core.retry.AsyncStreamingRetry] = None, + timeout: Optional[float] = None, + ) -> AsyncUploadOperation: + """Resumes an existing upload asynchronously, returning an AsyncUploadOperation. + + Args: + upload_url (str): Established upload session URL. + stream (Union[AsyncIterable[bytes], BinaryIO, bytes, Iterable[bytes]]): + 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. + retry (Optional[google.api_core.retry.AsyncStreamingRetry]): Optional + retry configuration (``AsyncStreamingRetry``) for chunk upload + requests. Use this to customize exponential backoff timing + between chunk retries or to supply a custom ``predicate`` for + API-specific transient errors. A custom ``predicate`` replaces + the default transient HTTP status check (HTTP 408, 429, 500, + 502, 503, and 504). Transport errors and protocol recovery + errors (HTTP 400, 412, 416, and ``MissingStatusHeaderError``) + always initiate server offset recovery. Terminal errors + (``DeadlineExceeded``, ``TransferStalledError``, + ``UploadCancelledError``, and ``UnseekableStreamError``) are + never retried. + timeout (Optional[float]): Optional per-attempt timeout ceiling in + seconds. + + Returns: + An AsyncUploadOperation handle representing the resumed transfer. + + Raises: + 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.") + actual_url = upload_url or self.upload_url + if not actual_url: + raise ValueError("An upload URL must be provided to resume.") + if stream is None: + raise ValueError("A data stream or payload must be provided to resume.") + + self._reset_transfer_state() + if chunk_size is not None: + self._state._chunk_size = chunk_size + + self._state._upload_url = actual_url + progress_queue: List[UploadProgress] = [] + reader_fn, computed_size, stream_obj = self._prepare_async_reader(stream, size) + + async def _run() -> AsyncGenerator[UploadProgress, None]: + try: + await self._recover(sess, stream_obj, progress_queue=progress_queue) + while progress_queue: + yield progress_queue.pop(0) + + async for item in self._transmit_all_chunks( + sess, + reader_fn, + computed_size, + progress_queue=progress_queue, + stream_obj=stream_obj, + retry=retry, + timeout=timeout, + ): + yield item + except BaseException as exc: + self._enrich_exception(exc) + raise + + return AsyncUploadOperation(progress_stream=_run(), session=self) + + def _prepare_async_reader( + self, + stream: Union[AsyncIterable[bytes], BinaryIO, bytes, Iterable[bytes]], + size: Optional[int], + ) -> Tuple[ + Callable[[int], Awaitable[bytes]], + Optional[int], + Optional[object], + ]: + """Creates an asynchronous byte reader and determines stream length. + + Args: + stream: Input payload (async iterable, binary stream, bytes, or iterable). + size: Explicit total size in bytes, if known. + + Returns: + Tuple of (async reader function, computed total size, underlying stream object). + + Raises: + TypeError: If the stream type is not supported. + """ + computed_size = size + + if isinstance(stream, (str, dict)): + raise TypeError(f"Unsupported stream type: {type(stream)}") + + if isinstance(stream, bytes): + bytes_io = io.BytesIO(stream) + if computed_size is None: + computed_size = len(stream) + + async def reader(n: int) -> bytes: + return bytes_io.read(n) + + return reader, computed_size, bytes_io + + read_fn = getattr(stream, "read", None) + if callable(read_fn): + if inspect.iscoroutinefunction(read_fn): + # Native async reader (e.g. asyncio.StreamReader) + async def reader(n: int) -> bytes: + return await read_fn(n) + + return reader, computed_size, stream + + # Synchronous binary stream (e.g. io.BytesIO or open file handle): + # offload blocking reads to worker thread via asyncio.to_thread. + if computed_size is None: + computed_size = _get_buffer_size(stream) + + tell_fn = getattr(stream, "tell", None) + if callable(tell_fn): + try: + self._start_stream_offset = tell_fn() + except (OSError, AttributeError): + self._start_stream_offset = 0 + + async def reader(n: int) -> bytes: + return await asyncio.to_thread(read_fn, n) + + return reader, computed_size, stream + + if hasattr(stream, "__aiter__"): + # Native AsyncIterable[bytes] + iterator = stream.__aiter__() + buffer = bytearray() + + async def reader(n: int) -> bytes: + while len(buffer) < n: + try: + chunk = await iterator.__anext__() + buffer.extend(chunk) + except StopAsyncIteration: + break + result = bytes(buffer[:n]) + del buffer[:n] + return result + + return reader, computed_size, None + + if isinstance(stream, Iterable): + # Synchronous Iterable[bytes]: offload to worker thread + iterator = iter(stream) + buffer = bytearray() + + def _next_chunk(): + try: + return next(iterator) + except StopIteration: + return None + + async def reader(n: int) -> bytes: + while len(buffer) < n: + chunk = await asyncio.to_thread(_next_chunk) + if chunk is None: + break + buffer.extend(chunk) + result = bytes(buffer[:n]) + del buffer[:n] + return result + + return reader, computed_size, None + + raise TypeError(f"Unsupported stream type: {type(stream)}") diff --git a/packages/google-api-core/google/api_core/resumable_transfer/upload_state.py b/packages/google-api-core/google/api_core/resumable_transfer/upload_state.py new file mode 100644 index 000000000000..df93a38e87e2 --- /dev/null +++ b/packages/google-api-core/google/api_core/resumable_transfer/upload_state.py @@ -0,0 +1,339 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sans-I/O Resumable Upload protocol state machine.""" + +import logging +from typing import Dict, Mapping, Optional, Sequence, Tuple, Union + +from google.api_core import exceptions +from google.api_core.resumable_transfer import common + +_LOGGER = logging.getLogger(__name__) + + +class _ProtocolState(object): + """Encapsulates the state and command formatting for Resumable Upload protocol.""" + + def __init__( + self, + upload_url: Optional[str] = None, + chunk_size: int = common.DEFAULT_CHUNK_SIZE, + ) -> None: + """Initializes the protocol state machine. + + Args: + upload_url: The initial service endpoint URL for starting a new + upload (stored statically as ``initial_url``), or an existing + session URL when resuming. + chunk_size: Desired chunk size in bytes. + """ + self._initial_url = upload_url or "" + self._chunk_size = chunk_size + self._upload_url: Optional[str] = upload_url + self._chunk_granularity: Optional[int] = None + self._bytes_uploaded = 0 + self._total_bytes: Optional[int] = None + self._finished = False + self._invalid = False + + @property + def initial_url(self) -> str: + """str: The static service endpoint URL used to initiate a new upload session. + + Unlike ``upload_url``, this value remains unchanged throughout the + lifecycle of the state machine and is used by ``build_start_request()`` + as the target endpoint for the ``start`` command. + """ + return self._initial_url + + @property + def upload_url(self) -> Optional[str]: + """Optional[str]: The established upload session URL, or ``None`` if not established. + + Unlike ``initial_url`` (which remains static as the start endpoint), + ``upload_url`` is populated with the unique session URL returned by the + server in the ``x-goog-upload-url`` header once the upload is initiated + (or updated via ``set_session_url()`` when resuming). It is used for + subsequent chunk transfer, query, and cancel requests, and can be saved + to resume an interrupted upload later. + """ + return self._upload_url + + @property + def bytes_uploaded(self) -> int: + """The confirmed number of bytes committed to the server.""" + return self._bytes_uploaded + + @property + def total_bytes(self) -> Optional[int]: + """The total payload size in bytes, or None if unknown.""" + return self._total_bytes + + @property + def finished(self) -> bool: + """Whether the upload has completed successfully.""" + return self._finished + + @property + def invalid(self) -> bool: + """Whether the upload session has encountered a terminal failure.""" + return self._invalid + + @property + def chunk_size(self) -> int: + """Block-aligned chunk size informed by server granularity.""" + if self._chunk_granularity: + return ( + (self._chunk_size + self._chunk_granularity - 1) + // self._chunk_granularity + ) * self._chunk_granularity + return self._chunk_size + + def build_start_request( + self, + body: Union[str, bytes] = "", + headers: Optional[Sequence[Tuple[str, str]]] = None, + content_type: Optional[str] = None, + size: Optional[int] = None, + ) -> Tuple[str, str, Dict[str, str], bytes]: + """Formats the HTTP start request. + + Args: + body: Initial metadata payload. + headers: Optional sequence of header tuples to include. + content_type: MIME type of the stream payload. + size: Total size of the stream in bytes, if known. + + Returns: + A tuple of (HTTP method, URL, headers dict, payload bytes). + + Raises: + ValueError: If the initial upload URL was not provided. + """ + if not self._initial_url: + raise ValueError("upload_url must be provided to start an upload.") + self._upload_url = None + self._chunk_granularity = None + self._bytes_uploaded = 0 + self._total_bytes = size + self._finished = False + self._invalid = False + req_headers: Dict[str, str] = {} + + if headers: + for k, v in headers: + key = k.decode("utf-8") if isinstance(k, bytes) else str(k) + val = v.decode("utf-8") if isinstance(v, bytes) else str(v) + req_headers[key] = val + + req_headers[common.HEADER_PROTOCOL] = common.PROTOCOL_RESUMABLE + req_headers[common.HEADER_COMMAND] = common._Command.START.value + + if content_type is not None: + req_headers[common.HEADER_CONTENT_TYPE] = content_type + if size is not None: + req_headers[common.HEADER_CONTENT_LENGTH] = str(size) + + payload = body.encode("utf-8") if isinstance(body, str) else body + return "POST", self._initial_url, req_headers, payload + + def process_start_response( + self, status_code: int, headers: Mapping[str, str] + ) -> str: + """Processes start response and extracts upload session URL. + + Args: + status_code: HTTP response status code. + headers: HTTP response headers. + + Returns: + The established resumable upload session URL. + + Raises: + exceptions.MissingStatusHeaderError: If status header is missing. + ValueError: If start response indicates failure or URL is missing. + """ + if status_code not in (200, 201): + self._invalid = True + raise ValueError(f"Start command failed with status {status_code}") + + headers_lower = {k.lower(): v for k, v in headers.items()} + status = headers_lower.get(common.HEADER_STATUS.lower()) + if not status: + raise exceptions.MissingStatusHeaderError( + f"Missing {common.HEADER_STATUS} header in start response" + ) + + upload_url = headers_lower.get(common.HEADER_URL.lower()) + if not upload_url: + self._invalid = True + raise ValueError(f"Server did not return {common.HEADER_URL} header") + + self._upload_url = upload_url + granularity = headers_lower.get(common.HEADER_CHUNK_GRANULARITY.lower()) + if granularity: + self._chunk_granularity = int(granularity) + + return self._upload_url + + def build_chunk_request( + self, + data: Union[bytes, memoryview], + is_last_chunk: bool, + content_type: Optional[str] = None, + ) -> Tuple[str, str, Dict[str, str], bytes]: + """Formats an upload chunk request. + + Args: + data: Chunk byte data or memoryview slice. + is_last_chunk: True if this chunk concludes the upload payload. + content_type: MIME type of the uploaded chunk data. + + Returns: + A tuple of (HTTP method, URL, headers dict, payload bytes). + + Raises: + ValueError: If upload session URL is not established. + """ + if not self._upload_url: + raise ValueError("Upload session URL not established.") + + command = ( + f"{common._Command.UPLOAD.value}, {common._Command.FINALIZE.value}" + if is_last_chunk + else common._Command.UPLOAD.value + ) + + headers = { + common.HEADER_COMMAND: command, + common.HEADER_OFFSET: str(self._bytes_uploaded), + } + if content_type: + headers["Content-Type"] = content_type + + payload = bytes(data) if isinstance(data, memoryview) else data + return "POST", self._upload_url, headers, payload + + def process_chunk_response( + self, status_code: int, headers: Mapping[str, str], chunk_bytes_sent: int + ) -> None: + """Processes upload chunk response and updates committed bytes. + + Args: + status_code: HTTP response status code. + headers: HTTP response headers. + chunk_bytes_sent: Byte length of the chunk sent in the request. + + Raises: + exceptions.MissingStatusHeaderError: If status header is missing from successful response. + exceptions.UploadCancelledError: If server indicates session was cancelled. + """ + if status_code not in (200, 201): + return + + headers_lower = {k.lower(): v for k, v in headers.items()} + status = headers_lower.get(common.HEADER_STATUS.lower()) + if not status: + raise exceptions.MissingStatusHeaderError( + f"Missing {common.HEADER_STATUS} header in chunk upload response" + ) + + if status == common._Status.ACTIVE.value: + self._bytes_uploaded += chunk_bytes_sent + elif status == common._Status.FINAL.value: + self._finished = True + self._bytes_uploaded += chunk_bytes_sent + elif status == common._Status.CANCELLED.value: + self._invalid = True + raise exceptions.UploadCancelledError( + "Upload session was cancelled by server" + ) + + def build_query_request(self) -> Tuple[str, str, Dict[str, str], bytes]: + """Formats the query request to discover server offset during recovery. + + Returns: + A tuple of (HTTP method, URL, headers dict, payload bytes). + + Raises: + ValueError: If upload session URL is not established. + """ + if not self._upload_url: + raise ValueError("Upload session URL not established.") + + headers = {common.HEADER_COMMAND: common._Command.QUERY.value} + return "POST", self._upload_url, headers, b"" + + def process_query_response( + self, status_code: int, headers: Mapping[str, str] + ) -> int: + """Processes query response and returns current server byte offset. + + Args: + status_code: HTTP response status code. + headers: HTTP response headers. + + Returns: + The current server byte offset. + + Raises: + ValueError: If query recovery indicates failure. + exceptions.UploadCancelledError: If server indicates session was cancelled. + """ + if status_code not in (200, 201): + self._invalid = True + raise ValueError(f"Query recovery failed with status {status_code}") + + headers_lower = {k.lower(): v for k, v in headers.items()} + status = headers_lower.get(common.HEADER_STATUS.lower()) + + if status == common._Status.ACTIVE.value: + received = int(headers_lower.get(common.HEADER_SIZE_RECEIVED.lower(), "0")) + self._bytes_uploaded = received + elif status == common._Status.FINAL.value: + self._finished = True + elif status == common._Status.CANCELLED.value: + self._invalid = True + raise exceptions.UploadCancelledError( + "Upload session was cancelled by server" + ) + + return self._bytes_uploaded + + def build_cancel_request(self) -> Tuple[str, str, Dict[str, str], bytes]: + """Formats the cancel request. + + Returns: + A tuple of (HTTP method, URL, headers dict, payload bytes). + + Raises: + ValueError: If upload session URL is not established. + """ + if not self._upload_url: + raise ValueError("Upload session URL not established.") + + headers = {common.HEADER_COMMAND: common._Command.CANCEL.value} + return "POST", self._upload_url, headers, b"" + + def process_cancel_response( + self, status_code: int, headers: Mapping[str, str] + ) -> None: + """Processes cancel response and marks session invalid. + + Args: + status_code: HTTP response status code. + headers: HTTP response headers. + """ + self._invalid = True diff --git a/packages/google-api-core/tests/asyncio/test_grpc_helpers_async.py b/packages/google-api-core/tests/asyncio/test_grpc_helpers_async.py index dcb09f18fea2..c90c6c7bceeb 100644 --- a/packages/google-api-core/tests/asyncio/test_grpc_helpers_async.py +++ b/packages/google-api-core/tests/asyncio/test_grpc_helpers_async.py @@ -17,7 +17,7 @@ from unittest.mock import AsyncMock # pragma: NO COVER # noqa: F401 except ImportError: # pragma: NO COVER import mock # type: ignore -import pytest # noqa: I202 +import pytest from ..helpers import warn_deprecated_credentials_file diff --git a/packages/google-api-core/tests/asyncio/test_rest_streaming_async.py b/packages/google-api-core/tests/asyncio/test_rest_streaming_async.py index 743a60fb05ab..61919aac5604 100644 --- a/packages/google-api-core/tests/asyncio/test_rest_streaming_async.py +++ b/packages/google-api-core/tests/asyncio/test_rest_streaming_async.py @@ -28,7 +28,7 @@ import mock # type: ignore import proto -import pytest # noqa: I202 +import pytest try: from google.auth.aio.transport import Response 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 new file mode 100644 index 000000000000..0645e362cb15 --- /dev/null +++ b/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py @@ -0,0 +1,2242 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Asynchronous tests for Resumable Upload protocol implementation.""" + +import asyncio +import datetime +import io +import json +from typing import ( + Any, + AsyncIterable, + AsyncIterator, + Dict, + List, + Mapping, + Optional, + Tuple, + Union, +) +from unittest import mock + +import pytest +from google.protobuf import empty_pb2 + +from google.api_core import exceptions +from google.api_core.resumable_transfer import ( + AsyncResumableUploadSession, + AsyncUploadOperation, + ProgressState, + ResumableUploadConfig, + TransferStalledError, + UnseekableStreamError, + UploadCancelledError, + UploadProgress, + common, + upload_async, +) +from tests.helpers import EchoResponse + +try: + import aiohttp # noqa: F401 + import google.auth.aio.transport # noqa: F401 + + GOOGLE_AUTH_AIO_INSTALLED = True +except ImportError: + GOOGLE_AUTH_AIO_INSTALLED = False + + +@pytest.fixture(autouse=True) +def check_async_rest_installed(request: pytest.FixtureRequest) -> None: + if request.node.name == "test_async_ensure_aiohttp_missing": + return + if not GOOGLE_AUTH_AIO_INSTALLED: + pytest.skip("Skipped because google-api-core[async_rest] is not installed") + + +class DummyResponse: + """Mock response class representing a deserialized protobuf message.""" + + def __init__(self, name: str, size: int) -> None: + """Initializes a DummyResponse. + + Args: + name: Resource name string. + size: Resource size in bytes. + """ + self.name = name + self.size = size + + @classmethod + def from_json(cls, data: Union[str, bytes]) -> "DummyResponse": + """Deserializes JSON payload into a DummyResponse instance. + + Args: + data: JSON byte string or text. + + Returns: + A DummyResponse instance. + """ + d = json.loads(data.decode("utf-8") if isinstance(data, bytes) else data) + return cls(name=d.get("name", ""), size=d.get("size", 0)) + + +class DummyAsyncResponse: + """Mock HTTP response conforming to aiohttp.ClientResponse interface.""" + + def __init__( + self, + status: int = 200, + headers: Optional[Mapping[str, str]] = None, + body: bytes = b"", + ) -> None: + """Initializes a DummyAsyncResponse. + + Args: + status: HTTP status code. + headers: HTTP response headers mapping. + body: Response payload bytes. + """ + self.status = status + self.headers = headers or {} + self._body = body + + async def read(self) -> bytes: + """Reads and returns response payload bytes. + + Returns: + Raw response payload bytes. + """ + return self._body + + async def __aenter__(self) -> "DummyAsyncResponse": + """Enters the asynchronous context manager. + + Returns: + The DummyAsyncResponse instance. + """ + return self + + async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + """Exits the asynchronous context manager. + + Args: + exc_type: Exception type if raised. + exc_val: Exception value if raised. + exc_tb: Exception traceback if raised. + """ + pass + + +class DummyAsyncSession: + """Mock asynchronous HTTP client session conforming to aiohttp.ClientSession interface.""" + + def __init__(self, responses: Optional[List[DummyAsyncResponse]] = None) -> None: + """Initializes a DummyAsyncSession. + + Args: + responses: Sequence of canned DummyAsyncResponse objects. + """ + self._responses: List[DummyAsyncResponse] = list(responses or []) + self.requests: List[Tuple[str, str, Dict[str, Any]]] = [] + + def request(self, method: str, url: str, **kwargs: Any) -> DummyAsyncResponse: + """Records the request and yields the next canned response. + + Args: + method: HTTP method verb. + url: Destination endpoint URL. + **kwargs: Additional request parameters. + + Returns: + Next canned DummyAsyncResponse. + """ + self.requests.append((method, url, kwargs)) + if not self._responses: + return DummyAsyncResponse(status=200, headers={}, body=b"") + return self._responses.pop(0) + + +class NonSeekableBytesIO(io.BytesIO): + """BytesIO stream simulation with seekable returning False.""" + + def seekable(self) -> bool: + """Reports whether the stream supports random access. + + Returns: + False unconditionally. + """ + return False + + +# ===================================================================== +# 1. Initialization and Configuration Tests +# ===================================================================== + + +def test_async_session_initialization_defaults() -> None: + """Validates default attribute values of an uninitiated async session.""" + session = AsyncResumableUploadSession(upload_url="https://api.example.com/start") + assert session.upload_url == "https://api.example.com/start" + assert session.chunk_size == common.DEFAULT_CHUNK_SIZE + assert session.response is None + assert session.bytes_uploaded == 0 + assert session.finished is False + + +def test_async_missing_transport_raises() -> None: + """Verifies that invoking session operations without a transport raises ValueError.""" + session = AsyncResumableUploadSession() + + with pytest.raises(ValueError, match="aiohttp.ClientSession"): + session.upload(stream=b"data") + + with pytest.raises(ValueError, match="aiohttp.ClientSession"): + session.resume(upload_url="https://upload.example.com", stream=b"data") + + +@pytest.mark.asyncio +async def test_async_cancel_missing_transport_raises() -> None: + """Verifies that cancel without a transport raises ValueError.""" + session = AsyncResumableUploadSession() + with pytest.raises(ValueError, match="aiohttp.ClientSession"): + await session.cancel() + + +def test_async_ensure_aiohttp_missing(monkeypatch: pytest.MonkeyPatch) -> None: + """Verifies that _ensure_aiohttp raises ImportError when aiohttp is unavailable.""" + monkeypatch.setattr(upload_async, "_HAS_AIOHTTP", False) + session = AsyncResumableUploadSession() + with pytest.raises(ImportError, match="google-api-core\\[async_rest\\]"): + session._ensure_aiohttp() + + +# ===================================================================== +# 2. Upload Execution and Operation Handle Tests +# ===================================================================== + + +@pytest.mark.asyncio +async def test_async_upload_direct_execution() -> None: + """Verifies single-chunk upload execution with protobuf deserialization.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"name": "async_file.txt", "size": 10}', + ) + + async_transport = DummyAsyncSession([start_resp, chunk_resp]) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=DummyResponse, + transport=async_transport, + ) + + result = await session.upload(stream=b"0123456789", request_body='{"name": "test"}') + + assert isinstance(result, DummyResponse) + assert result.name == "async_file.txt" + assert result.size == 10 + assert session.finished is True + assert session.bytes_uploaded == 10 + assert session.upload_url == "https://upload.example.com/resumable-async" + + +@pytest.mark.asyncio +async def test_async_upload_multi_chunk_operation_handle() -> None: + """Verifies multi-chunk upload dispatching and AsyncUploadOperation property handles.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk1_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "active"}, + body=b"", + ) + chunk2_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"name": "multi.txt", "size": 8}', + ) + + async_transport = DummyAsyncSession([start_resp, chunk1_resp, chunk2_resp]) + config = ResumableUploadConfig(chunk_size=4) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + response_type=DummyResponse, + transport=async_transport, + ) + + upload_op = session.upload(stream=b"12345678") + assert isinstance(upload_op, AsyncUploadOperation) + assert upload_op.chunk_size == 4 + + result = await upload_op + assert isinstance(result, DummyResponse) + assert result.name == "multi.txt" + assert upload_op.response == result + assert upload_op.bytes_uploaded == 8 + assert upload_op.upload_url == "https://upload.example.com/resumable-async" + + # Validate commands dispatched in request history + assert len(async_transport.requests) == 3 + # Start request + assert async_transport.requests[0][2]["headers"]["X-Goog-Upload-Command"] == "start" + # Chunk 1 request + assert ( + async_transport.requests[1][2]["headers"]["X-Goog-Upload-Command"] == "upload" + ) + assert async_transport.requests[1][2]["headers"]["X-Goog-Upload-Offset"] == "0" + # Chunk 2 request (last chunk concludes transfer) + assert ( + async_transport.requests[2][2]["headers"]["X-Goog-Upload-Command"] + == "upload, finalize" + ) + assert async_transport.requests[2][2]["headers"]["X-Goog-Upload-Offset"] == "4" + + +@pytest.mark.asyncio +async def test_async_upload_progress_tracking() -> None: + """Verifies that progress stream yields snapshots matching transmission milestones.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk1_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "active"}, + body=b"", + ) + chunk2_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"name": "progress.txt", "size": 8}', + ) + + async_transport = DummyAsyncSession([start_resp, chunk1_resp, chunk2_resp]) + config = ResumableUploadConfig(chunk_size=4) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + response_type=DummyResponse, + transport=async_transport, + ) + + upload_op = session.upload(stream=b"12345678") + progress_list: List[UploadProgress] = [] + async for p in upload_op.progress(): + progress_list.append(p) + + final_resp = await upload_op + assert isinstance(final_resp, DummyResponse) + assert len(progress_list) == 3 + assert progress_list[0].state == ProgressState.STARTED + assert progress_list[1].state == ProgressState.UPLOADING + assert progress_list[1].bytes_uploaded == 4 + assert progress_list[2].state == ProgressState.FINALIZED + assert progress_list[2].bytes_uploaded == 8 + + +# ===================================================================== +# 3. Stream Input Types Tests +# ===================================================================== + + +@pytest.mark.asyncio +async def test_async_stream_types_async_iterable() -> None: + """Verifies upload compatibility with an asynchronous generator stream.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"name": "async_gen.txt", "size": 6}', + ) + + async_transport = DummyAsyncSession([start_resp, chunk_resp]) + + async def async_generator() -> AsyncIterator[bytes]: + yield b"abc" + yield b"def" + + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=DummyResponse, + transport=async_transport, + ) + + resp = await session.upload(stream=async_generator()) + assert resp.name == "async_gen.txt" + assert session.bytes_uploaded == 6 + + +@pytest.mark.asyncio +async def test_async_stream_types_binary_io() -> None: + """Verifies upload compatibility with a seekable BinaryIO stream.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"name": "bytes_io.txt", "size": 5}', + ) + + async_transport = DummyAsyncSession([start_resp, chunk_resp]) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=DummyResponse, + transport=async_transport, + ) + + stream = io.BytesIO(b"hello") + resp = await session.upload(stream=stream) + assert resp.name == "bytes_io.txt" + assert session.bytes_uploaded == 5 + + +@pytest.mark.asyncio +async def test_async_stream_types_sync_iterable() -> None: + """Verifies upload compatibility with a synchronous iterable of byte chunks.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"name": "iterable.txt", "size": 6}', + ) + + async_transport = DummyAsyncSession([start_resp, chunk_resp]) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=DummyResponse, + transport=async_transport, + ) + + resp = await session.upload(stream=[b"foo", b"bar"]) + assert resp.name == "iterable.txt" + assert session.bytes_uploaded == 6 + + +def test_async_stream_types_unsupported_raises() -> None: + """Verifies that passing an unsupported stream type raises TypeError.""" + session = AsyncResumableUploadSession(transport=DummyAsyncSession()) + prepare_reader = getattr(session, "_prepare_async_reader") + with pytest.raises(TypeError, match="Unsupported stream type"): + prepare_reader(stream=12345, size=10) + + +# ===================================================================== +# 4. Resume and Offset Recovery Tests +# ===================================================================== + + +@pytest.mark.asyncio +async def test_async_resume_success() -> None: + """Verifies resuming an existing upload by querying server offset.""" + query_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "5", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"name": "resumed_async.txt", "size": 10}', + ) + + async_transport = DummyAsyncSession([query_resp, chunk_resp]) + session = AsyncResumableUploadSession( + response_type=DummyResponse, + transport=async_transport, + ) + + upload_op = session.resume( + upload_url="https://upload.example.com/resumable-async", + stream=b"0123456789", + ) + resp = await upload_op + assert resp.name == "resumed_async.txt" + assert session.bytes_uploaded == 10 + assert session.finished is True + + # First request was query + assert async_transport.requests[0][2]["headers"]["X-Goog-Upload-Command"] == "query" + # Second request was remaining chunk starting from offset 5 + assert async_transport.requests[1][2]["headers"]["X-Goog-Upload-Offset"] == "5" + + +def test_async_resume_missing_arguments_raises() -> None: + """Verifies that resume raises ValueError when upload_url or stream is missing.""" + session = AsyncResumableUploadSession(transport=DummyAsyncSession()) + with pytest.raises(ValueError, match="An upload URL must be provided to resume"): + session.resume(upload_url="", stream=b"data") + + with pytest.raises( + ValueError, match="A data stream or payload must be provided to resume" + ): + session.resume( + upload_url="https://upload.example.com/resumable-async", + stream=None, # type: ignore[arg-type] + ) + + +@pytest.mark.asyncio +async def test_async_resume_recovery_unseekable_stream_raises() -> None: + """Verifies that UnseekableStreamError is raised if server offset cannot be rewound.""" + query_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "100", + }, + body=b"", + ) + + async_transport = DummyAsyncSession([query_resp]) + session = AsyncResumableUploadSession( + response_type=DummyResponse, + transport=async_transport, + ) + + stream = NonSeekableBytesIO(b"some content") + upload_op = session.resume( + upload_url="https://upload.example.com/resumable-async", + stream=stream, + ) + + with pytest.raises(UnseekableStreamError) as exc_info: + await upload_op + + assert exc_info.value.upload_url == "https://upload.example.com/resumable-async" + + +@pytest.mark.asyncio +async def test_async_resume_recovery_seekable_stream() -> None: + """Verifies that seekable streams are rewound to committed server offset.""" + query_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "3", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"name": "seekable.txt", "size": 6}', + ) + + async_transport = DummyAsyncSession([query_resp, chunk_resp]) + session = AsyncResumableUploadSession( + response_type=DummyResponse, + transport=async_transport, + ) + + stream = io.BytesIO(b"abcdef") + upload_op = session.resume( + upload_url="https://upload.example.com/resumable-async", + stream=stream, + ) + resp = await upload_op + assert resp.name == "seekable.txt" + assert session.bytes_uploaded == 6 + + +# ===================================================================== +# 5. Cancellation Tests +# ===================================================================== + + +@pytest.mark.asyncio +async def test_async_cancel_success() -> None: + """Verifies client-initiated cancellation marks session state invalid.""" + cancel_resp = DummyAsyncResponse(status=200, headers={}, body=b"") + async_transport = DummyAsyncSession([cancel_resp]) + + session = AsyncResumableUploadSession( + upload_url="https://upload.example.com/resumable-async", + transport=async_transport, + ) + await session.cancel() + assert session._state.invalid is True + assert ( + async_transport.requests[0][2]["headers"]["X-Goog-Upload-Command"] == "cancel" + ) + + +@pytest.mark.asyncio +async def test_async_server_cancelled_raises_error() -> None: + """Verifies that server returning cancelled status raises UploadCancelledError.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "cancelled"}, + body=b"", + ) + + async_transport = DummyAsyncSession([start_resp, chunk_resp]) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + transport=async_transport, + ) + + with pytest.raises(UploadCancelledError, match="cancelled by server"): + await session.upload(stream=b"12345") + + +# ===================================================================== +# 6. Retry and Error Handling Tests +# ===================================================================== + + +@pytest.mark.asyncio +async def test_async_retry_transient_http_errors() -> None: + """Verifies transparent retries on transient HTTP status codes (503 Service Unavailable).""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_503 = DummyAsyncResponse(status=503, headers={}, body=b"Service Unavailable") + chunk_success = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"name": "retried.txt", "size": 5}', + ) + + async_transport = DummyAsyncSession([start_resp, chunk_503, chunk_success]) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=DummyResponse, + transport=async_transport, + ) + + with mock.patch("asyncio.sleep", new_callable=mock.AsyncMock): + resp = await session.upload(stream=b"hello") + + assert resp.name == "retried.txt" + assert session.finished is True + + +@pytest.mark.asyncio +async def test_async_non_retryable_error_raises() -> None: + """Verifies that non-retryable errors (e.g. 404 Not Found) terminate immediately.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_404 = DummyAsyncResponse(status=404, headers={}, body=b"Not Found") + + async_transport = DummyAsyncSession([start_resp, chunk_404]) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + transport=async_transport, + ) + + with pytest.raises(exceptions.NotFound): + await session.upload(stream=b"test") + + +@pytest.mark.asyncio +async def test_async_recoverable_status_code_triggers_recovery() -> None: + """Verifies that recoverable error status codes trigger query and offset reconciliation.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + # Chunk 1 returns Category 2 recoverable error (412 Precondition Failed) + chunk_precondition_failed = DummyAsyncResponse( + status=412, headers={}, body=b"Precondition Failed" + ) + # Recovery query returns confirmed committed offset 0 + query_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "0", + }, + body=b"", + ) + # Retransmission succeeds + chunk_success = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"name": "recovered.txt", "size": 5}', + ) + + async_transport = DummyAsyncSession( + [start_resp, chunk_precondition_failed, query_resp, chunk_success] + ) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=DummyResponse, + transport=async_transport, + ) + + with mock.patch("asyncio.sleep", new_callable=mock.AsyncMock): + resp = await session.upload(stream=b"hello") + + assert resp.name == "recovered.txt" + assert session.bytes_uploaded == 5 + + +@pytest.mark.asyncio +async def test_async_missing_status_header_triggers_recovery() -> None: + """Verifies that missing status header triggers query recovery.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + # Successful HTTP status but missing X-Goog-Upload-Status header + chunk_missing_hdr = DummyAsyncResponse(status=200, headers={}, body=b"") + # Recovery query + query_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "0", + }, + body=b"", + ) + chunk_success = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"name": "header_recovered.txt", "size": 5}', + ) + + async_transport = DummyAsyncSession( + [start_resp, chunk_missing_hdr, query_resp, chunk_success] + ) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=DummyResponse, + transport=async_transport, + ) + + with mock.patch("asyncio.sleep", new_callable=mock.AsyncMock): + resp = await session.upload(stream=b"hello") + + assert resp.name == "header_recovered.txt" + + +# ===================================================================== +# 7. Stall Control and Deadline Tests +# ===================================================================== + + +@pytest.mark.asyncio +async def test_async_stall_timeout_raises_transfer_stalled_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verifies that transfer stalling beyond timeout threshold raises TransferStalledError.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"name": "slow.txt", "size": 10}', + ) + + async_transport = DummyAsyncSession([start_resp, chunk_resp]) + # Expect 100 bytes/sec, timeout 1 second + config = ResumableUploadConfig(stall_minimum_rate=100, stall_timeout=1.0) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=async_transport, + ) + + # Simulate elapsed time 10.0 seconds during 10-byte upload (rate = 1 byte/s < 100) + clock_vals = iter([0.0, 10.0, 10.0]) + monkeypatch.setattr(upload_async, "_monotonic_clock", lambda: next(clock_vals)) + + with pytest.raises(TransferStalledError, match="Upload stalled"): + await session.upload(stream=b"0123456789") + + # config has stall_minimum_rate=100 B/s, stall_timeout=1.0s. + # 100 bytes / 100 B/s -> expected_sec = 1.0s. + # Taking 1.2s gives current_lag = 1.2 - 1.0 = 0.2s. + # At clock = 10.0s, _stall_timeout_started is backdated by current_lag + # (10.0 - 0.2 = 9.8s) so elapsed stall time is 0.2s (< 1.0s stall_timeout). + clock_vals3 = iter([10.0, 11.0]) + monkeypatch.setattr(upload_async, "_monotonic_clock", lambda: next(clock_vals3)) + session3 = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + ) + session3._update_stall_control(100, 1.2) + assert session3._aggregate_lag == pytest.approx(0.2) + assert session3._stall_timeout_started == pytest.approx(9.8) + + # At clock = 11.0s, elapsed time since _stall_timeout_started (9.8s) is + # 11.0 - 9.8 = 1.2s >= stall_timeout (1.0s), raising TransferStalledError. + with pytest.raises(TransferStalledError, match="Upload stalled"): + session3._update_stall_control(100, 1.2) + + # With aggregate_lag = 0.4s (0.2s + 0.2s), next_chunk_timeout is + # 1.0 - 0.4 + 1.0 = 1.6s. timeout_override=0.5s is smaller, so 0.5s wins. + assert session3._compute_chunk_timeout(100, timeout_override=0.5) == pytest.approx( + 0.5 + ) + session_no_stall = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=ResumableUploadConfig(stall_minimum_rate=0, stall_timeout=0), + ) + assert session_no_stall._compute_chunk_timeout( + 100, timeout_override=15.0 + ) == pytest.approx(15.0) + + +@pytest.mark.asyncio +async def test_async_deadline_exceeded() -> None: + """Verifies that exceeding the configured upload deadline raises DeadlineExceeded.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + async_transport = DummyAsyncSession([start_resp]) + past_deadline = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta( + seconds=10 + ) + config = ResumableUploadConfig(deadline=past_deadline) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=async_transport, + ) + + with pytest.raises(exceptions.DeadlineExceeded): + await session.upload(stream=b"data") + + +# ===================================================================== +# 8. Response Deserialization Types Tests +# ===================================================================== + + +@pytest.mark.asyncio +async def test_async_response_type_proto_message() -> None: + """Verifies that a proto.Message type parses final response bytes.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"content": "proto_async_payload"}', + ) + + async_transport = DummyAsyncSession([start_resp, chunk_resp]) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=EchoResponse, + transport=async_transport, + ) + + result = await session.upload(stream=b"data") + assert isinstance(result, EchoResponse) + assert result.content == "proto_async_payload" + + +@pytest.mark.asyncio +async def test_async_response_type_protobuf_message() -> None: + """Verifies that a google.protobuf.message.Message type parses final response bytes.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b"{}", + ) + + async_transport = DummyAsyncSession([start_resp, chunk_resp]) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=empty_pb2.Empty, + transport=async_transport, + ) + + result = await session.upload(stream=b"data") + assert isinstance(result, empty_pb2.Empty) + + +@pytest.mark.asyncio +async def test_async_response_type_callable() -> None: + """Verifies that a custom callable deserializer parses final response body bytes.""" + if not GOOGLE_AUTH_AIO_INSTALLED: + pytest.skip("Skipped because google-api-core[async_rest] is not installed") + + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b"parsed:hello", + ) + + async_transport = DummyAsyncSession([start_resp, chunk_resp]) + + def custom_parser(raw: bytes) -> str: + return raw.decode("utf-8").upper() + + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=custom_parser, + transport=async_transport, + ) + + result = await session.upload(stream=b"data") + assert result == "PARSED:HELLO" + + +@pytest.mark.asyncio +async def test_async_response_type_raw_bytes() -> None: + """Verifies that raw bytes are returned when response_type is not configured.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b"raw-bytes-output", + ) + + async_transport = DummyAsyncSession([start_resp, chunk_resp]) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + transport=async_transport, + ) + + result = await session.upload(stream=b"data") + assert result == b"raw-bytes-output" + + +@pytest.mark.asyncio +async def test_async_operation_error_propagation_in_progress() -> None: + """Verifies that background task errors propagate through progress queue iteration.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=403, + headers={}, + body=b"Permission Denied", + ) + + async_transport = DummyAsyncSession([start_resp, chunk_resp]) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + transport=async_transport, + ) + + upload_op = session.upload(stream=b"data") + with pytest.raises(exceptions.Forbidden): + async for _ in upload_op.progress(): + pass + + with pytest.raises(exceptions.Forbidden): + await upload_op + + +@pytest.mark.parametrize("invalid_stream", ["invalid_string", {"key": "value"}, 12345]) +def test_async_upload_rejects_invalid_stream_types(invalid_stream: Any) -> None: + """Verifies that str, dict, and non-stream objects raise TypeError synchronously on upload().""" + async_transport = DummyAsyncSession([]) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + transport=async_transport, + ) + with pytest.raises(TypeError, match="Unsupported stream type"): + session.upload(stream=invalid_stream) + + +@pytest.mark.parametrize("invalid_stream", ["invalid_string", {"key": "value"}, 12345]) +def test_async_resume_rejects_invalid_stream_types(invalid_stream: Any) -> None: + """Verifies that str, dict, and non-stream objects raise TypeError synchronously on resume().""" + async_transport = DummyAsyncSession([]) + session = AsyncResumableUploadSession( + transport=async_transport, + ) + with pytest.raises(TypeError, match="Unsupported stream type"): + session.resume( + upload_url="https://upload.example.com/resumable-async", + stream=invalid_stream, + ) + + +def test_async_enrich_exception() -> None: + """Verifies that _enrich_exception attaches upload_url and chunk_size.""" + session = AsyncResumableUploadSession( + upload_url="https://upload.example.com/resumable-async", + ) + exc = RuntimeError("test error") + session._enrich_exception(exc) + assert getattr(exc, "upload_url") == "https://upload.example.com/resumable-async" + + +def test_async_notify_progress_branches() -> None: + """Verifies progress notification queue capture.""" + session = AsyncResumableUploadSession( + upload_url=None, + ) + # When upload_url is None + session._notify_progress(common.ProgressState.UPLOADING) + + # When upload_url is established on state + session._state._upload_url = "https://upload.example.com/resumable-async" + # Call with queue=None + session._notify_progress(common.ProgressState.UPLOADING, queue=None) + + q: List[UploadProgress] = [] + session._notify_progress(common.ProgressState.UPLOADING, queue=q) + assert len(q) == 1 + + +def test_async_deadline_handling_and_start_timeout() -> None: + """Verifies deadline remaining calculations and start timeout calculation.""" + # Past deadline raises DeadlineExceeded + past = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=10) + config = ResumableUploadConfig(deadline=past) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + ) + with pytest.raises(exceptions.DeadlineExceeded): + session._get_deadline_remaining() + + # Naive future deadline is localized to UTC + future_naive = datetime.datetime.now() + datetime.timedelta(hours=1) + config2 = ResumableUploadConfig(deadline=future_naive) + session2 = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config2, + ) + rem = session2._get_deadline_remaining() + assert rem is not None and rem > 0 + t = session2._get_start_timeout() + assert t > 0 + + +@pytest.mark.asyncio +async def test_async_retry_branches() -> None: + """Verifies retry predicate and policy resolution branches in upload_async.""" + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + ) + + pred_start = session._get_retry_predicate(is_start=True) + pred_transfer = session._get_retry_predicate(is_start=False) + + assert pred_start(exceptions.DeadlineExceeded("deadline")) is False + assert pred_start(exceptions.TransferStalledError("stalled")) is False + assert pred_start(exceptions.UploadCancelledError("cancelled")) is False + assert pred_start(exceptions.UnseekableStreamError("unseekable")) is False + assert pred_start(exceptions.MissingStatusHeaderError("missing")) is False + assert pred_transfer(exceptions.MissingStatusHeaderError("missing")) is True + assert pred_start(aiohttp.ClientError("network error")) is True + assert pred_start(exceptions.from_http_status(503, "Service Unavailable")) is True + assert pred_start(exceptions.from_http_status(400, "Bad Request")) is False + assert pred_start(exceptions.from_http_status(412, "Precondition Failed")) is False + assert ( + pred_start(exceptions.from_http_status(416, "Range Not Satisfiable")) is False + ) + assert pred_start(exceptions.from_http_status(409, "Conflict")) is False + assert pred_start(RuntimeError("runtime")) is False + + # Category 1 and Category 2 status codes are retryable during transfer + assert pred_transfer(exceptions.from_http_status(500, "Internal Error")) is True + assert pred_transfer(exceptions.from_http_status(400, "Bad Request")) is True + assert ( + pred_transfer(exceptions.from_http_status(412, "Precondition Failed")) is True + ) + assert ( + pred_transfer(exceptions.from_http_status(416, "Range Not Satisfiable")) is True + ) + assert pred_transfer(exceptions.from_http_status(409, "Conflict")) is False + + # Default retry resolution + default_unary = session._get_async_retry() + assert isinstance(default_unary, google.api_core.retry.AsyncRetry) + + default_stream = session._get_async_streaming_retry() + assert isinstance(default_stream, google.api_core.retry.AsyncStreamingRetry) + assert default_stream.timeout is None + + class CustomApiError(Exception): + """Example API-specific transient exception provided by a caller.""" + + # ------------------------------------------------------------------------- + # Scenario 1: User provides a custom predicate to retry an API-specific error + # ------------------------------------------------------------------------- + custom_unary = google.api_core.retry.AsyncRetry( + initial=0.5, + predicate=lambda exc: isinstance(exc, CustomApiError), + ) + session_unary = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + start_retry=custom_unary, + ) + resolved_unary = session_unary._get_async_retry() + + assert isinstance(resolved_unary, google.api_core.retry.AsyncRetry) + assert resolved_unary._initial == 0.5 + # User's custom exception is retried, while unrelated errors are not + assert resolved_unary._predicate(CustomApiError("rate limit")) is True + assert resolved_unary._predicate(ValueError("invalid input")) is False + # Terminal errors must always return False regardless of custom predicate + assert ( + resolved_unary._predicate(exceptions.DeadlineExceeded("deadline expired")) + is False + ) + + # ------------------------------------------------------------------------- + # Scenario 2: Custom AsyncStreamingRetry preserves Category 2 recovery + # ------------------------------------------------------------------------- + custom_stream = google.api_core.retry.AsyncStreamingRetry( + initial=0.5, + predicate=google.api_core.retry.if_exception_type(CustomApiError), + ) + converted_stream = session_unary._get_async_streaming_retry( + retry_override=custom_stream + ) + assert isinstance(converted_stream, google.api_core.retry.AsyncStreamingRetry) + assert converted_stream._initial == 0.5 + assert converted_stream._predicate(CustomApiError("rate limit")) is True + # Category 2 recovery (400, 412, 416, MissingStatusHeaderError) is preserved during chunk transfer + missing_header_error = exceptions.MissingStatusHeaderError( + "Missing X-Goog-Upload-Status" + ) + assert converted_stream._predicate(missing_header_error) is True + assert ( + converted_stream._predicate( + exceptions.from_http_status(412, "Precondition Failed") + ) + is True + ) + + # ------------------------------------------------------------------------- + # Scenario 3: Restrictive custom predicate still preserves Category 2 recovery + # ------------------------------------------------------------------------- + restrictive_stream = google.api_core.retry.AsyncStreamingRetry( + initial=0.25, + predicate=lambda exc: False, + ) + resolved_stream = session._get_async_streaming_retry( + retry_override=restrictive_stream + ) + assert isinstance(resolved_stream, google.api_core.retry.AsyncStreamingRetry) + assert resolved_stream._initial == 0.25 + # Category 2 errors (400, 412, 416, MissingStatusHeaderError) return True so + # server offset synchronization is preserved + assert ( + resolved_stream._predicate(exceptions.from_http_status(400, "Bad Request")) + is True + ) + assert ( + resolved_stream._predicate( + exceptions.from_http_status(412, "Precondition Failed") + ) + is True + ) + assert ( + resolved_stream._predicate( + exceptions.from_http_status(416, "Range Not Satisfiable") + ) + is True + ) + assert resolved_stream._predicate(missing_header_error) is True + # Unretriable HTTP status codes (such as 409 Conflict) and non-protocol errors + # follow the predicate and return False + assert ( + resolved_stream._predicate(exceptions.from_http_status(409, "Conflict")) + is False + ) + bad_gateway_error = exceptions.from_http_status(502, "Bad Gateway") + assert resolved_stream._predicate(bad_gateway_error) is False + assert resolved_stream._predicate(RuntimeError("unexpected crash")) is False + + +def test_async_transport_missing_errors() -> None: + """Verifies ValueError when transport is missing from upload, resume, and cancel.""" + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + ) + with pytest.raises( + ValueError, match="An aiohttp.ClientSession transport must be provided" + ): + session.upload(stream=b"data") + + with pytest.raises( + ValueError, match="An aiohttp.ClientSession transport must be provided" + ): + session.resume(upload_url="https://upload.example.com/123", stream=b"data") + + +@pytest.mark.asyncio +async def test_async_cancel_missing_transport_and_error() -> None: + """Verifies cancel method with missing transport and server error.""" + session = AsyncResumableUploadSession( + upload_url="https://upload.example.com/123", + ) + with pytest.raises( + ValueError, match="An aiohttp.ClientSession transport must be provided" + ): + await session.cancel() + + err_resp = DummyAsyncResponse(status=500, headers={}, body=b"Cancel Error") + sess_transport = DummyAsyncSession([err_resp]) + session2 = AsyncResumableUploadSession( + upload_url="https://upload.example.com/123", + transport=sess_transport, + ) + with pytest.raises(exceptions.GoogleAPICallError): + await session2.cancel() + + +@pytest.mark.asyncio +async def test_async_prepare_async_reader_types() -> None: + """Verifies async reader preparation for native async reader, tell error, and iterables.""" + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + ) + + class AsyncReader(AsyncIterable[bytes]): # Inherit to satisfy mypy + async def read(self, n: int) -> bytes: + return b"chunk" + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b"chunk" + + reader_fn, size, obj = session._prepare_async_reader(AsyncReader(), None) + chunk = await reader_fn(5) + assert chunk == b"chunk" + + # Sync stream whose tell() raises OSError + class TellFailingStream(io.BytesIO): + def tell(self) -> int: + raise OSError("tell error") + + stream = TellFailingStream(b"data") + reader_fn2, size2, obj2 = session._prepare_async_reader(stream, None) + chunk2 = await reader_fn2(4) + assert chunk2 == b"data" + assert session._start_stream_offset == 0 + + # Sync Iterable[bytes] + reader_fn3, size3, obj3 = session._prepare_async_reader([b"part1", b"part2"], None) + chunk3 = await reader_fn3(10) + assert chunk3 == b"part1part2" + + +@pytest.mark.asyncio +async def test_async_initiate_and_recover_failures() -> None: + """Verifies initiate and recover error handling when server returns error codes.""" + err_resp = DummyAsyncResponse(status=400, headers={}, body=b"Bad Request") + sess_transport = DummyAsyncSession([err_resp]) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + transport=sess_transport, + ) + with pytest.raises(exceptions.BadRequest): + await session._initiate(transport=sess_transport) + + err_resp2 = DummyAsyncResponse(status=400, headers={}, body=b"Query Failed") + sess_transport2 = DummyAsyncSession([err_resp2]) + session2 = AsyncResumableUploadSession( + transport=sess_transport2, + ) + session2._state._upload_url = "https://upload.example.com/123" + with pytest.raises(exceptions.BadRequest): + await session2._recover(sess_transport2) + + +@pytest.mark.asyncio +async def test_async_recover_stream_errors() -> None: + """Verifies UnseekableStreamError during async recovery.""" + query_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "active", "X-Goog-Upload-Size-Received": "10"}, + body=b"", + ) + sess_transport = DummyAsyncSession([query_resp]) + session = AsyncResumableUploadSession( + transport=sess_transport, + ) + session._state._upload_url = "https://upload.example.com/123" + + # Stream whose seekable() returns False + unseekable = mock.Mock() + unseekable.seekable.return_value = False + with pytest.raises(UnseekableStreamError, match="Stream is not seekable"): + await session._recover(sess_transport, stream_obj=unseekable) + + # Stream whose seek() raises OSError + sess_transport2 = DummyAsyncSession([query_resp]) + session2 = AsyncResumableUploadSession( + transport=sess_transport2, + ) + session2._state._upload_url = "https://upload.example.com/123" + failing_seek = mock.Mock() + failing_seek.seekable.return_value = True + failing_seek.seek.side_effect = OSError("Seek error") + with pytest.raises(UnseekableStreamError, match="Failed to seek stream"): + await session2._recover(sess_transport2, stream_obj=failing_seek) + + +@pytest.mark.asyncio +async def test_async_upload_with_timeout_and_deadline() -> None: + future_deadline = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta( + seconds=60 + ) + config = ResumableUploadConfig(deadline=future_deadline) + + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + resp = DummyAsyncResponse( + status=200, headers={"X-Goog-Upload-Status": "final"}, body=b"{}" + ) + sess_transport = DummyAsyncSession([start_resp, resp]) + + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=sess_transport, + ) + session._state._upload_url = "https://upload.example.com/resumable-async" + + # Run the upload + res = await session.upload(stream=b"data", timeout=30.0) + assert res == b"{}" + + +@pytest.mark.asyncio +async def test_async_transmit_chunk_timeout_errors() -> None: + # 1. TimeoutError raises TransferStalledError when remaining is > 0 + future_deadline = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta( + seconds=60 + ) + config = ResumableUploadConfig(deadline=future_deadline) + + class TimeoutAsyncSession: + def __init__(self): + self.calls = 0 + + def request(self, *args, **kwargs): + self.calls += 1 + if self.calls == 1: + + class StartContext: + async def __aenter__(self): + class Resp: + status = 200 + headers = { + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + } + + async def read(self): + return b"" + + return Resp() + + async def __aexit__(self, exc_type, exc, tb): + pass + + return StartContext() + else: + + class TimeoutContext: + async def __aenter__(self): + raise asyncio.TimeoutError("timeout") + + async def __aexit__(self, exc_type, exc, tb): + pass + + return TimeoutContext() + + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=TimeoutAsyncSession(), + ) + session._state._upload_url = "https://upload.example.com/resumable-async" + no_retry = google.api_core.retry.AsyncStreamingRetry( + predicate=lambda e: False, timeout=0.001 + ) + + with pytest.raises(exceptions.TransferStalledError): + await session.upload(stream=b"data", retry=no_retry) + + # 2. TimeoutError raises DeadlineExceeded when deadline expires + config2 = ResumableUploadConfig(deadline=future_deadline) + + session2 = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config2, + transport=TimeoutAsyncSession(), + ) + session2._state._upload_url = "https://upload.example.com/resumable-async" + + with mock.patch.object( + session2, + "_get_deadline_remaining", + side_effect=[5.0, 5.0, exceptions.DeadlineExceeded("Deadline exceeded")], + ): + with pytest.raises(exceptions.DeadlineExceeded): + await session2.upload(stream=b"data", retry=no_retry) + + # 3. Non-timeout RetryError re-raises RetryError rather than TransferStalledError + class ConnectionErrorAsyncSession: + def __init__(self): + self.calls = 0 + + def request(self, *args, **kwargs): + self.calls += 1 + if self.calls == 1: + + class StartContext: + async def __aenter__(self): + class Resp: + status = 200 + headers = { + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + } + + async def read(self): + return b"" + + return Resp() + + async def __aexit__(self, exc_type, exc, tb): + pass + + return StartContext() + else: + + class ErrorContext: + async def __aenter__(self): + raise aiohttp.ClientConnectionError("Connection dropped") + + async def __aexit__(self, exc_type, exc, tb): + pass + + return ErrorContext() + + session3 = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=ConnectionErrorAsyncSession(), + ) + session3._state._upload_url = "https://upload.example.com/resumable-async" + + with pytest.raises(exceptions.RetryError): + await session3.upload(stream=b"data", retry=no_retry) + + +@pytest.mark.asyncio +async def test_async_upload_multiple_chunks_async_iterable() -> None: + # chunk_size = 3, payload = b"012345" (6 bytes) + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk1_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "active"}, + body=b"", + ) + chunk2_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b"{}", + ) + sess_transport = DummyAsyncSession([start_resp, chunk1_resp, chunk2_resp]) + + async def async_gen(): + yield b"012" + yield b"345" + + config = ResumableUploadConfig(chunk_size=3) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=sess_transport, + ) + session._state._upload_url = "https://upload.example.com/resumable-async" + + res = await session.upload(stream=async_gen()) + assert res == b"{}" + + +@pytest.mark.asyncio +async def test_async_prepare_async_reader_additional_branches() -> None: + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + ) + + # Inherit from concrete io.BytesIO so mypy recognizes these test streams as + # valid BinaryIO instances natively. + class LocalNoTellStream(io.BytesIO): + """Simulates a stream that has read() but lacks getbuffer() and tell(). + + Overrides __getattribute__ to raise AttributeError for "tell" and + "getbuffer" so that hasattr(stream, "tell") and hasattr(stream, "getbuffer") + evaluate to False at runtime while remaining mypy-compliant. + """ + + def __getattribute__(self, name: str) -> Any: + if name in ("tell", "getbuffer"): + raise AttributeError(f"no {name}") + return super().__getattribute__(name) + + class LocalCustomReadStream(io.BytesIO): + """Simulates a stream that lacks getbuffer() and where tell() raises OSError. + + Used to verify that _prepare_async_reader gracefully catches OSError + when attempting to record the starting stream offset via tell(). + """ + + def __getattribute__(self, name: str) -> Any: + if name == "getbuffer": + raise AttributeError("no getbuffer") + return super().__getattribute__(name) + + def tell(self) -> int: + raise OSError("tell failed") + + # 1. bytes stream with explicit size + reader_fn1, size1, obj1 = session._prepare_async_reader(b"data", size=4) + assert size1 == 4 + + # 2. NoTellStream: read but no tell + stream2 = LocalNoTellStream(b"") + reader_fn2, size2, obj2 = session._prepare_async_reader(stream2, size=None) + assert size2 is None + + # 3. CustomReadStream: read, but no getbuffer and tell raises OSError + stream3 = LocalCustomReadStream(b"hello") + reader_fn3, size3, obj3 = session._prepare_async_reader(stream3, size=None) + assert size3 is None + + # 4. BinaryIO stream with explicit size (covers computed_size is not None branch) + reader_fn4, size4, obj4 = session._prepare_async_reader( + io.BytesIO(b"hello"), size=5 + ) + assert size4 == 5 + + +@pytest.mark.asyncio +async def test_async_upload_empty_stream() -> None: + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b"{}", + ) + async_transport = DummyAsyncSession([start_resp, chunk_resp]) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + transport=async_transport, + ) + res = await session.upload(stream=b"") + assert res == b"{}" + + +@pytest.mark.asyncio +async def test_async_upload_no_stall_config() -> None: + config = ResumableUploadConfig(stall_minimum_rate=0, stall_timeout=0) + + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + resp = DummyAsyncResponse( + status=200, headers={"X-Goog-Upload-Status": "final"}, body=b"{}" + ) + sess_transport = DummyAsyncSession([start_resp, resp]) + + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=sess_transport, + ) + res = await session.upload(stream=b"data") + assert res == b"{}" + + +@pytest.mark.asyncio +async def test_async_transmit_chunk_timeout_errors_no_stall() -> None: + # 1. TimeoutError raises TransferStalledError when remaining is > 0 and no stall control + future_deadline = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta( + seconds=60 + ) + config = ResumableUploadConfig( + deadline=future_deadline, stall_minimum_rate=0, stall_timeout=0 + ) + + class TimeoutAsyncSession: + def __init__(self): + self.calls = 0 + + def request(self, *args, **kwargs): + self.calls += 1 + if self.calls == 1: + + class StartContext: + async def __aenter__(self): + class Resp: + status = 200 + headers = { + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + } + + async def read(self): + return b"" + + return Resp() + + async def __aexit__(self, exc_type, exc, tb): + pass + + return StartContext() + else: + + class TimeoutContext: + async def __aenter__(self): + raise asyncio.TimeoutError("timeout") + + async def __aexit__(self, exc_type, exc, tb): + pass + + return TimeoutContext() + + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=TimeoutAsyncSession(), + ) + session._state._upload_url = "https://upload.example.com/resumable-async" + no_retry = google.api_core.retry.AsyncStreamingRetry( + predicate=lambda e: False, timeout=0.001 + ) + + with pytest.raises(exceptions.TransferStalledError): + await session.upload(stream=b"data", retry=no_retry) + + # 2. TimeoutError raises DeadlineExceeded when deadline expires and no stall control + config2 = ResumableUploadConfig( + deadline=future_deadline, stall_minimum_rate=0, stall_timeout=0 + ) + + session2 = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config2, + transport=TimeoutAsyncSession(), + ) + session2._state._upload_url = "https://upload.example.com/resumable-async" + + with mock.patch.object( + session2, + "_get_deadline_remaining", + side_effect=[5.0, 5.0, exceptions.DeadlineExceeded("Deadline exceeded")], + ): + with pytest.raises(exceptions.DeadlineExceeded): + await session2.upload(stream=b"data", retry=no_retry) + + +@pytest.mark.asyncio +async def test_async_per_attempt_timeout_retries_before_stall_timeout( + monkeypatch, +) -> None: + """Verifies that hitting per_attempt_timeout (5s) in async uploads retries via recovery rather + than prematurely raising TransferStalledError when stall_timeout (120s) has not elapsed.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + query_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "0", + }, + body=b"", + ) + final_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"status": "completed"}', + ) + + class RecoverAfterTimeoutSession: + def __init__(self): + self.calls = 0 + + def request(self, *args, **kwargs): + self.calls += 1 + if self.calls == 1: + return start_resp + if self.calls == 2: + + class TimeoutContext: + async def __aenter__(self): + raise aiohttp.ServerTimeoutError("Per-attempt 5s timeout") + + async def __aexit__(self, exc_type, exc, tb): + pass + + return TimeoutContext() + if self.calls == 3: + return query_resp + return final_resp + + transport = RecoverAfterTimeoutSession() + config = ResumableUploadConfig( + stall_minimum_rate=1024, + stall_timeout=120.0, + ) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=transport, + ) + + clock_vals = iter([0.0, 5.0, 5.0, 5.0, 6.0, 6.0]) + monkeypatch.setattr(upload_async, "_monotonic_clock", lambda: next(clock_vals)) + + result = await session.upload( + stream=b"test data", + retry=google.api_core.retry.AsyncStreamingRetry(initial=0.01, maximum=0.01), + ) + assert result == b'{"status": "completed"}' + assert transport.calls == 4 + + +@pytest.mark.asyncio +async def test_async_recover_buffered_chunk_out_of_bounds() -> None: + session = AsyncResumableUploadSession() + session._state._upload_url = "https://upload.example.com/resumable-async" + session._buffered_chunk = memoryview(b"data") + session._buffered_chunk_offset = 0 + + query_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "active", "X-Goog-Upload-Size-Received": "10"}, + body=b"", + ) + sess_transport = DummyAsyncSession([query_resp]) + + unseekable = mock.Mock() + unseekable.seekable.return_value = False + + with pytest.raises(UnseekableStreamError): + await session._recover(sess_transport, stream_obj=unseekable) + + assert session._buffered_chunk is None + + +@pytest.mark.asyncio +async def test_async_recover_stream_obj_none() -> None: + session = AsyncResumableUploadSession() + session._state._upload_url = "https://upload.example.com/resumable-async" + session._buffered_chunk = None + + query_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "active", "X-Goog-Upload-Size-Received": "10"}, + body=b"", + ) + sess_transport = DummyAsyncSession([query_resp]) + + with pytest.raises(UnseekableStreamError, match="precedes active buffer"): + await session._recover(sess_transport, stream_obj=None) + + +@pytest.mark.asyncio +async def test_async_upload_already_finished_raises_value_error() -> None: + session = AsyncResumableUploadSession() + session._state._upload_url = "https://upload.example.com/resumable-async" + + async def mark_finished(*args, **kwargs): + session._state._finished = True + + with mock.patch.object( + session, "_initiate", new_callable=mock.AsyncMock, side_effect=mark_finished + ): + sess_transport = DummyAsyncSession([]) + with pytest.raises( + ValueError, match="Upload completed without receiving a final response" + ): + await session.upload(stream=b"data", transport=sess_transport) + + +@pytest.mark.asyncio +async def test_async_resume_already_finished_raises_value_error() -> None: + session = AsyncResumableUploadSession() + session._state._upload_url = "https://upload.example.com/resumable-async" + + async def mark_finished(*args, **kwargs): + session._state._finished = True + + sess_transport = DummyAsyncSession([]) + with mock.patch.object( + session, "_recover", new_callable=mock.AsyncMock, side_effect=mark_finished + ): + op = session.resume( + upload_url="https://upload.example.com/resumable-async", + stream=b"data", + chunk_size=1024, + transport=sess_transport, + ) + with pytest.raises( + ValueError, + match="Upload completed without receiving a final response", + ): + await op + + +@pytest.mark.asyncio +async def test_async_partial_chunk_recovery_does_not_prematurely_finalize() -> None: + """Ensure that retrying a partially committed chunk (len < chunk_size) does not prematurely finalize. + + When a 4-byte chunk (b"0123") partially succeeds (server commits 2 bytes) + and is retried, ensure that the remaining 2 bytes (b"23") are sent with + "upload" rather than "upload, finalize" so the remaining payload (b"45") + is not dropped. + """ + server_received_bytes = bytearray() + + class StatefulAsyncTransport: + def __init__(self) -> None: + self.call_count = 0 + + def request( + self, + method: str, + url: str, + data: Any = None, + headers: Optional[Mapping[str, str]] = None, + **kwargs: Any, + ) -> DummyAsyncResponse: + self.call_count += 1 + cmd = headers.get("X-Goog-Upload-Command", "") if headers else "" + + # Request 1: start session + if self.call_count == 1: + assert cmd == "start" + return DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + + # Request 2: initial Chunk 1 (b"0123") -> server commits partial 2 bytes (b"01"), then fails 503 + if self.call_count == 2: + assert cmd == "upload" + assert bytes(data) == b"0123" + server_received_bytes.extend(b"01") + return DummyAsyncResponse( + status=503, + headers={}, + body=b"Service Unavailable", + ) + + # Request 3: recovery query -> server reports 2 committed bytes + if self.call_count == 3: + assert cmd == "query" + return DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": str(len(server_received_bytes)), + }, + body=b"", + ) + + # Subsequent upload requests (Request 4: remaining b"23", Request 5: final b"45") + if data: + server_received_bytes.extend(bytes(data)) + + if "finalize" in cmd: + return DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b"{}", + ) + return DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "active"}, + body=b"", + ) + + transport = StatefulAsyncTransport() + config = ResumableUploadConfig( + chunk_size=4, + ) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=transport, + ) + + await session.upload( + stream=b"012345", + retry=google.api_core.retry.AsyncStreamingRetry( + predicate=lambda exc: True, initial=0.001 + ), + ) + + # Verify no data loss occurred: server must receive all 6 bytes (b"012345"), not truncated b"0123" + assert bytes(server_received_bytes) == b"012345" + + +@pytest.mark.asyncio +async def test_async_upload_progress_cancellation() -> None: + """Verifies that cancelling a task iterating over progress() propagates CancelledError cleanly.""" + slow_event = asyncio.Event() + + class HangingTransport: + def request(self, *args: Any, **kwargs: Any) -> Any: + class HangingCtx: + async def __aenter__(self) -> Any: + # Block indefinitely to simulate an in-flight HTTP request. + await slow_event.wait() + return DummyAsyncResponse(status=200, headers={}, body=b"") + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + pass + + return HangingCtx() + + session_cancel = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + transport=HangingTransport(), + ) + op_cancel = session_cancel.upload(stream=b"data") + + async def consume_progress() -> None: + async for _ in op_cancel.progress(): + pass + + # Start iterating over progress() in a background task, allow it to reach + # the blocking request, and then cancel the task. + progress_task = asyncio.create_task(consume_progress()) + await asyncio.sleep(0.01) + progress_task.cancel() + with pytest.raises(asyncio.CancelledError): + await progress_task + + +@pytest.mark.asyncio +async def test_async_upload_progress_base_exception() -> None: + """Verifies that a BaseException raised during progress() is captured and propagated.""" + + class CustomBaseException(BaseException): + pass + + class BaseExceptionTransport: + def request(self, *args: Any, **kwargs: Any) -> Any: + class BaseExceptionCtx: + async def __aenter__(self) -> Any: + # Raise a direct BaseException subclass to verify non-Exception + # errors mark the operation as consumed and propagate out of progress(). + raise CustomBaseException("fatal error") + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + pass + + return BaseExceptionCtx() + + session_base_exc = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + transport=BaseExceptionTransport(), + ) + op_base_exc = session_base_exc.upload(stream=b"data") + with pytest.raises(CustomBaseException, match="fatal error"): + async for _ in op_base_exc.progress(): + pass + assert op_base_exc._consumed is True + assert isinstance(op_base_exc._exception, CustomBaseException) + + +@pytest.mark.asyncio +async def test_async_upload_progress_repeated_iteration_after_completion() -> None: + """Verifies that iterating over progress() again after stream completion yields no items.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-async", + }, + body=b"", + ) + final_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b"{}", + ) + session_ok = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + transport=DummyAsyncSession([start_resp, final_resp]), + ) + op_ok = session_ok.upload(stream=b"data") + + # First pass drains the entire progress stream (STARTED and FINALIZED). + first_pass = [p async for p in op_ok] + assert len(first_pass) == 2 + assert op_ok._consumed is True + + # Subsequent iteration over progress() sees _consumed=True and returns immediately. + second_pass = [p async for p in op_ok.progress()] + assert second_pass == [] + + +@pytest.mark.asyncio +async def test_async_upload_partial_progress_iteration_then_await() -> None: + """Verifies that breaking out of progress iteration early allows awaiting the remaining upload.""" + start_partial = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-partial", + }, + body=b"", + ) + chunk1_partial = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "active"}, + body=b"", + ) + chunk2_partial = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b'{"name": "partial_then_await.txt", "size": 8}', + ) + # Configure a 2-chunk upload (8 bytes total with 4-byte chunk_size), which yields + # 3 progress notifications in total: STARTED, UPLOADING (after chunk 1), and FINALIZED (after chunk 2). + session_partial = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=ResumableUploadConfig(chunk_size=4), + response_type=DummyResponse, + transport=DummyAsyncSession([start_partial, chunk1_partial, chunk2_partial]), + ) + op_partial = session_partial.upload(stream=b"01234567") + + # Consume only the first two progress events (STARTED and first UPLOADING) and break early. + seen_states = [] + async for p in op_partial: + seen_states.append(p.state) + if len(seen_states) == 2: + break + + # The stream has not reached the end yet, so _consumed remains False. + assert op_partial._consumed is False + + # Awaiting the operation handle resumes draining the remaining chunks from _progress_stream + # until completion, marks _consumed=True, and returns the deserialized response. + final_result = await op_partial + assert op_partial._consumed is True + assert isinstance(final_result, DummyResponse) + assert final_result.name == "partial_then_await.txt" + assert final_result.size == 8 + + +@pytest.mark.asyncio +async def test_async_method_override_arguments() -> None: + """Verifies content_type and on_progress overrides on initiate, upload, and resume.""" + start_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/123", + }, + body=b"", + ) + chunk_resp = DummyAsyncResponse( + status=200, + headers={"X-Goog-Upload-Status": "final"}, + body=b"{}", + ) + + # 1. initiate with content_type override + session1 = AsyncResumableUploadSession(upload_url="https://api.example.com/start") + await session1._initiate( + transport=DummyAsyncSession([start_resp]), content_type="text/plain" + ) + assert session1._content_type == "text/plain" + + # 2. upload with content_type override + session2 = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + transport=DummyAsyncSession([start_resp, chunk_resp]), + ) + await session2.upload( + stream=b"data", + content_type="text/csv", + ) + assert session2._content_type == "text/csv" + + # 3. resume with chunk_size override + query_resp = DummyAsyncResponse( + status=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "0", + }, + body=b"", + ) + session3 = AsyncResumableUploadSession( + transport=DummyAsyncSession([query_resp, chunk_resp]) + ) + await session3.resume( + upload_url="https://upload.example.com/123", + stream=b"data", + ) + + +def test_async_retry_predicate_includes_asyncio_timeout_error() -> None: + """Verifies that asyncio.TimeoutError is treated as retryable by _get_retry_predicate.""" + session = AsyncResumableUploadSession(upload_url="https://api.example.com/start") + predicate = session._get_retry_predicate() + assert predicate(asyncio.TimeoutError()) is True diff --git a/packages/google-api-core/tests/helpers.py b/packages/google-api-core/tests/helpers.py index 279ebe9108b9..c630adb8b11c 100644 --- a/packages/google-api-core/tests/helpers.py +++ b/packages/google-api-core/tests/helpers.py @@ -19,7 +19,7 @@ from typing import List import proto -import pytest # noqa: I202 +import pytest from google.protobuf import duration_pb2, timestamp_pb2 from google.protobuf.json_format import MessageToJson diff --git a/packages/google-api-core/tests/unit/test_resumable_transfer.py b/packages/google-api-core/tests/unit/test_resumable_transfer.py new file mode 100644 index 000000000000..df4505f439fd --- /dev/null +++ b/packages/google-api-core/tests/unit/test_resumable_transfer.py @@ -0,0 +1,1944 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import datetime +import io +from typing import Union +from unittest import mock + +import pytest +import requests +from google.protobuf import empty_pb2 + +import google.api_core.retry +from google.api_core import exceptions +from google.api_core.resumable_transfer import ( + DEFAULT_CHUNK_SIZE, + MissingStatusHeaderError, + ProgressState, + ResumableUploadConfig, + ResumableUploadSession, + TransferStalledError, + UnseekableStreamError, + UploadCancelledError, + UploadProgress, + common, + upload, + upload_state, +) +from tests.helpers import EchoResponse + + +class DummyResponse: + """Mock response class representing a deserialized protobuf message.""" + + def __init__(self, name: str, size: int) -> None: + """Initializes a DummyResponse. + + Args: + name: Resource name string. + size: Resource size in bytes. + """ + self.name = name + self.size = size + + @classmethod + def from_json(cls, data: Union[str, bytes]) -> "DummyResponse": + """Deserializes JSON payload into a DummyResponse instance. + + Args: + data: JSON byte string or text. + + Returns: + A DummyResponse instance. + """ + import json + + d = json.loads(data.decode("utf-8") if isinstance(data, bytes) else data) + return cls(name=d.get("name", ""), size=d.get("size", 0)) + + +# ===================================================================== +# 1. Common Constants and Error Types +# ===================================================================== + + +def test_common_constants(): + assert DEFAULT_CHUNK_SIZE == 10 * 1024 * 1024 + assert common.HEADER_PROTOCOL == "X-Goog-Upload-Protocol" + assert common.HEADER_COMMAND == "X-Goog-Upload-Command" + assert common.HEADER_STATUS == "X-Goog-Upload-Status" + assert common.HEADER_URL == "X-Goog-Upload-URL" + assert common.HEADER_OFFSET == "X-Goog-Upload-Offset" + assert common.HEADER_SIZE_RECEIVED == "X-Goog-Upload-Size-Received" + assert common.PROTOCOL_RESUMABLE == "resumable" + + assert common._Command.START == "start" + assert common._Command.UPLOAD == "upload" + assert common._Command.FINALIZE == "finalize" + assert common._Command.QUERY == "query" + assert common._Command.CANCEL == "cancel" + + assert common._Status.ACTIVE == "active" + assert common._Status.FINAL == "final" + assert common._Status.CANCELLED == "cancelled" + + assert ProgressState.STARTED == "started" + assert ProgressState.UPLOADING == "uploading" + assert ProgressState.RECOVERING == "recovering" + assert ProgressState.OFFSET_RECEIVED == "offset received" + assert ProgressState.FINALIZED == "finalized" + + +def test_upload_progress_dataclass(): + prog = UploadProgress( + upload_url="https://upload.example.com/session123", + chunk_size=1024, + bytes_uploaded=512, + total_bytes=2048, + state=ProgressState.UPLOADING, + ) + assert prog.upload_url == "https://upload.example.com/session123" + assert prog.chunk_size == 1024 + assert prog.bytes_uploaded == 512 + assert prog.total_bytes == 2048 + assert prog.state == ProgressState.UPLOADING + + +def test_exception_hierarchy(): + assert issubclass(TransferStalledError, exceptions.GoogleAPICallError) + assert issubclass(UnseekableStreamError, exceptions.GoogleAPICallError) + assert issubclass(UploadCancelledError, exceptions.GoogleAPICallError) + assert issubclass(MissingStatusHeaderError, exceptions.GoogleAPICallError) + assert exceptions.TransferStalledError is TransferStalledError + assert exceptions.UnseekableStreamError is UnseekableStreamError + assert exceptions.UploadCancelledError is UploadCancelledError + assert exceptions.MissingStatusHeaderError is MissingStatusHeaderError + + +# ===================================================================== +# 2. Pure Sans-I/O State Machine (upload_state.py) +# ===================================================================== + + +def test_protocol_state_start_request(): + state = upload_state._ProtocolState(upload_url="https://api.example.com/start") + method, url, headers, payload = state.build_start_request( + body='{"name": "test"}', + headers=[("X-Custom", "val")], + content_type="text/plain", + size=1000, + ) + + assert method == "POST" + assert url == "https://api.example.com/start" + assert headers["X-Goog-Upload-Protocol"] == "resumable" + assert headers["X-Goog-Upload-Command"] == "start" + assert headers["X-Goog-Upload-Header-Content-Type"] == "text/plain" + assert headers["X-Goog-Upload-Header-Content-Length"] == "1000" + assert headers["X-Custom"] == "val" + assert payload == b'{"name": "test"}' + + +def test_protocol_state_process_start_response(): + state = upload_state._ProtocolState(upload_url="https://api.example.com/start") + headers = { + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-id", + "X-Goog-Upload-Chunk-Granularity": "262144", + } + url = state.process_start_response(200, headers) + assert url == "https://upload.example.com/resumable-id" + assert state.upload_url == "https://upload.example.com/resumable-id" + assert state._chunk_granularity == 262144 + + +def test_protocol_state_start_response_missing_status(): + state = upload_state._ProtocolState(upload_url="https://api.example.com/start") + headers = {"X-Goog-Upload-URL": "https://upload.example.com/resumable-id"} + with pytest.raises(MissingStatusHeaderError): + state.process_start_response(200, headers) + + +def test_protocol_state_start_response_missing_url(): + state = upload_state._ProtocolState(upload_url="https://api.example.com/start") + headers = {"X-Goog-Upload-Status": "active"} + with pytest.raises(ValueError, match="Server did not return"): + state.process_start_response(200, headers) + + +def test_protocol_state_granularity_alignment(): + state = upload_state._ProtocolState(chunk_size=500) + assert state.chunk_size == 500 + state._chunk_granularity = 256 + # 500 rounded up to multiple of 256 is 512 + assert state.chunk_size == 512 + + +def test_protocol_state_chunk_request_and_response(): + state = upload_state._ProtocolState(upload_url="https://upload.example.com/session") + + # First chunk: not last + method, url, headers, payload = state.build_chunk_request( + data=b"0123456789", is_last_chunk=False + ) + assert headers["X-Goog-Upload-Command"] == "upload" + assert headers["X-Goog-Upload-Offset"] == "0" + assert payload == b"0123456789" + + state.process_chunk_response(200, {"X-Goog-Upload-Status": "active"}, 10) + assert state.bytes_uploaded == 10 + assert not state.finished + + # Second chunk: last chunk + method, url, headers, payload = state.build_chunk_request( + data=b"abcdef", is_last_chunk=True, content_type="text/plain" + ) + assert headers["X-Goog-Upload-Command"] == "upload, finalize" + assert headers["X-Goog-Upload-Offset"] == "10" + assert headers["Content-Type"] == "text/plain" + + state.process_chunk_response(200, {"X-Goog-Upload-Status": "final"}, 6) + assert state.bytes_uploaded == 16 + assert state.finished + + +def test_protocol_state_chunk_missing_status_header(): + state = upload_state._ProtocolState(upload_url="https://upload.example.com/session") + with pytest.raises(MissingStatusHeaderError): + state.process_chunk_response(200, {}, 10) + + +def test_protocol_state_chunk_cancelled_status(): + state = upload_state._ProtocolState(upload_url="https://upload.example.com/session") + with pytest.raises(UploadCancelledError): + state.process_chunk_response(200, {"X-Goog-Upload-Status": "cancelled"}, 10) + assert state.invalid + + +def test_protocol_state_query_and_cancel(): + state = upload_state._ProtocolState(upload_url="https://upload.example.com/session") + method, url, headers, payload = state.build_query_request() + assert headers["X-Goog-Upload-Command"] == "query" + + received = state.process_query_response( + 200, {"X-Goog-Upload-Status": "active", "X-Goog-Upload-Size-Received": "1024"} + ) + assert received == 1024 + assert state.bytes_uploaded == 1024 + + # query with unknown status + state2 = upload_state._ProtocolState( + upload_url="https://upload.example.com/session" + ) + received2 = state2.process_query_response(200, {"X-Goog-Upload-Status": "unknown"}) + assert received2 == 0 + + method, url, headers, payload = state.build_cancel_request() + assert headers["X-Goog-Upload-Command"] == "cancel" + state.process_cancel_response(200, {}) + assert state.invalid + + +# ===================================================================== +# 3. ResumableUploadConfig Sensible Defaults +# ===================================================================== + + +def test_resumable_upload_config_defaults(): + config = ResumableUploadConfig() + assert config.chunk_size == 10 * 1024 * 1024 + assert config.stall_minimum_rate == 64 * 1024 + assert config.stall_timeout == 120.0 + assert config.headers is None + assert config.deadline is None + + +def test_resumable_upload_config_headers(): + config1 = ResumableUploadConfig( + headers={"X-Test": "1"}, + ) + assert config1.start_headers == [("X-Test", "1")] + + config2 = ResumableUploadConfig( + headers=[("X-Test", "2")], + ) + assert config2.start_headers == [("X-Test", "2")] + + +# ===================================================================== +# 4. Synchronous ResumableUploadSession (upload.py) +# ===================================================================== + + +def test_sync_upload_direct_execution(): + session_transport = mock.create_autospec(requests.Session, instance=True) + + # 1. Start response + start_resp = mock.Mock() + start_resp.ok = True + start_resp.status_code = 200 + start_resp.headers = { + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-123", + } + + # 2. Chunk response + chunk_resp = mock.Mock() + chunk_resp.ok = True + chunk_resp.status_code = 200 + chunk_resp.headers = {"X-Goog-Upload-Status": "final"} + chunk_resp.content = b'{"name": "done.txt", "size": 11}' + + session_transport.request.side_effect = [start_resp, chunk_resp] + + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + transport=session_transport, + response_type=DummyResponse, + ) + + payload = b"Hello world" + result = session.upload(stream=payload, request_body='{"name": "test"}') + + assert isinstance(result, DummyResponse) + assert result.name == "done.txt" + assert result.size == 11 + assert session.finished is True + assert session.bytes_uploaded == 11 + assert session.upload_url == "https://upload.example.com/resumable-123" + assert session.response == result + + +def test_sync_upload_iterative_progress(): + session_transport = mock.create_autospec(requests.Session, instance=True) + + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-123", + }, + ) + chunk1_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "active"}, + ) + chunk2_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b'{"name": "stream.txt", "size": 8}', + ) + + session_transport.request.side_effect = [start_resp, chunk1_resp, chunk2_resp] + + config = ResumableUploadConfig(chunk_size=4) + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=session_transport, + response_type=DummyResponse, + ) + + progress_events = list(session.iter_upload(stream=b"12345678")) + assert len(progress_events) == 3 + assert progress_events[0].state == ProgressState.STARTED + assert progress_events[1].state == ProgressState.UPLOADING + assert progress_events[1].bytes_uploaded == 4 + assert progress_events[2].state == ProgressState.FINALIZED + assert progress_events[2].bytes_uploaded == 8 + + assert session.response.name == "stream.txt" + assert session.response.size == 8 + + +def test_sync_resume(): + session_transport = mock.create_autospec(requests.Session, instance=True) + + # 1. Query response returns offset 5 + query_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "5", + }, + ) + # 2. Remaining chunk response + chunk_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b'{"name": "resumed.txt", "size": 10}', + ) + session_transport.request.side_effect = [query_resp, chunk_resp] + + session = ResumableUploadSession(response_type=DummyResponse) + + stream = io.BytesIO(b"0123456789") + resp = session.resume( + upload_url="https://upload.example.com/resumable-123", + stream=stream, + transport=session_transport, + ) + + assert isinstance(resp, DummyResponse) + assert resp.name == "resumed.txt" + assert session.bytes_uploaded == 10 + assert session.finished is True + + +def test_sync_iter_resume(): + """Verifies streaming progress during upload resumption.""" + session_transport = mock.create_autospec(requests.Session, instance=True) + + query_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "5", + }, + ) + chunk_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b'{"name": "iter_resumed.txt", "size": 10}', + ) + session_transport.request.side_effect = [query_resp, chunk_resp] + + session = ResumableUploadSession(response_type=DummyResponse) + + stream = io.BytesIO(b"0123456789") + progress_list = list( + session.iter_resume( + upload_url="https://upload.example.com/resumable-123", + stream=stream, + transport=session_transport, + ) + ) + + assert session.response.name == "iter_resumed.txt" + assert session.bytes_uploaded == 10 + assert session.finished is True + assert len(progress_list) == 2 + assert progress_list[0].state == ProgressState.OFFSET_RECEIVED + assert progress_list[1].state == ProgressState.FINALIZED + + +def test_sync_recoverable_status_code_triggers_offset_recovery(): + session_transport = mock.create_autospec(requests.Session, instance=True) + + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-123", + }, + ) + # First chunk upload fails with 400 (recoverable Category 2) + err400_resp = mock.Mock( + ok=False, + status_code=400, + headers={}, + ) + err400_resp.json.return_value = {"error": {"message": "Bad Request", "details": []}} + err400_resp.text = '{"error": {"message": "Bad Request"}}' + err400_resp.content = err400_resp.text.encode("utf-8") + # Recovery query returns offset 0 + query_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "0", + }, + ) + # Retry chunk upload succeeds + success_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b'{"name": "recovered.txt", "size": 5}', + ) + + session_transport.request.side_effect = [ + start_resp, + err400_resp, + query_resp, + success_resp, + ] + + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=DummyResponse, + transport=session_transport, + ) + + resp = session.upload(stream=b"12345") + assert resp.name == "recovered.txt" + assert session.bytes_uploaded == 5 + + +def test_sync_exceptions_carry_upload_url_and_chunk_size(): + session_transport = mock.create_autospec(requests.Session, instance=True) + + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-123", + }, + ) + # Fatal 403 error on chunk upload + err403_resp = mock.Mock( + ok=False, + status_code=403, + headers={}, + ) + err403_resp.json.return_value = {"error": {"message": "Forbidden", "details": []}} + err403_resp.text = '{"error": {"message": "Forbidden"}}' + err403_resp.content = err403_resp.text.encode("utf-8") + session_transport.request.side_effect = [start_resp, err403_resp] + + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + transport=session_transport, + ) + + with pytest.raises(exceptions.GoogleAPICallError) as exc_info: + session.upload(stream=b"data") + + err = exc_info.value + assert err.upload_url == "https://upload.example.com/resumable-123" + assert err.chunk_size == session.chunk_size + + +def test_sync_unseekable_stream_error_on_preceding_offset(): + session_transport = mock.create_autospec(requests.Session, instance=True) + + query_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "50", + }, + ) + session_transport.request.side_effect = [query_resp] + + session = ResumableUploadSession( + config=ResumableUploadConfig(), + transport=session_transport, + ) + + # Mock an unseekable stream + unseekable = mock.Mock(spec=io.RawIOBase) + unseekable.seekable.return_value = False + + with pytest.raises(UnseekableStreamError) as exc_info: + session.resume( + upload_url="https://upload.example.com/resumable-123", + stream=unseekable, + transport=session_transport, + ) + + assert exc_info.value.upload_url == "https://upload.example.com/resumable-123" + + +def test_sync_stall_timeout(monkeypatch): + session_transport = mock.create_autospec(requests.Session, instance=True) + + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-123", + }, + ) + session_transport.request.side_effect = [ + start_resp, + requests.exceptions.Timeout("Read timed out"), + ] + + # Configure stall control with 0.1s timeout + config = ResumableUploadConfig( + stall_minimum_rate=1024, + stall_timeout=0.1, + ) + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=session_transport, + ) + + clock_vals = iter([0.0, 0.5, 0.5]) + monkeypatch.setattr(upload, "_monotonic_clock", lambda: next(clock_vals)) + + with pytest.raises(TransferStalledError) as exc_info: + session.upload(stream=b"test data") + + assert exc_info.value.upload_url == "https://upload.example.com/resumable-123" + + +def test_sync_cancel(): + session_transport = mock.create_autospec(requests.Session, instance=True) + cancel_resp = mock.Mock(ok=True, status_code=200, headers={}) + session_transport.request.return_value = cancel_resp + + session = ResumableUploadSession( + upload_url="https://upload.example.com/resumable-123", + transport=session_transport, + ) + session.cancel() + assert session._state.invalid is True + + +def test_sync_response_type_proto_message(): + session_transport = mock.create_autospec(requests.Session, instance=True) + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-123", + }, + ) + chunk_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b'{"content": "proto_payload"}', + ) + session_transport.request.side_effect = [start_resp, chunk_resp] + + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=EchoResponse, + transport=session_transport, + ) + resp = session.upload(stream=b"payload") + assert isinstance(resp, EchoResponse) + assert resp.content == "proto_payload" + + +def test_sync_response_type_protobuf_message(): + session_transport = mock.create_autospec(requests.Session, instance=True) + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-123", + }, + ) + chunk_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b"{}", + ) + session_transport.request.side_effect = [start_resp, chunk_resp] + + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=empty_pb2.Empty, + transport=session_transport, + ) + resp = session.upload(stream=b"payload") + assert isinstance(resp, empty_pb2.Empty) + + +def test_sync_response_type_callable(): + session_transport = mock.create_autospec(requests.Session, instance=True) + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-123", + }, + ) + chunk_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b"custom_payload", + ) + session_transport.request.side_effect = [start_resp, chunk_resp] + + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=lambda c: c.decode("utf-8").upper(), + transport=session_transport, + ) + resp = session.upload(stream=b"payload") + assert resp == "CUSTOM_PAYLOAD" + + +def test_sync_response_type_raw_response(): + session_transport = mock.create_autospec(requests.Session, instance=True) + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-123", + }, + ) + chunk_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b"raw_content", + ) + session_transport.request.side_effect = [start_resp, chunk_resp] + + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + response_type=None, + transport=session_transport, + ) + resp = session.upload(stream=b"payload") + assert resp == b"raw_content" + + +def test_sync_retry_predicate_allows_timeout_with_stall_control(): + config = ResumableUploadConfig(stall_minimum_rate=1024, stall_timeout=1.0) + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + ) + predicate = session._get_retry_predicate() + assert predicate(requests.exceptions.Timeout("Read timed out")) is True + + +@pytest.mark.parametrize("invalid_stream", ["invalid_string", {"key": "value"}, 12345]) +def test_sync_upload_rejects_invalid_stream_types(invalid_stream): + session_transport = mock.create_autospec(requests.Session, instance=True) + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + transport=session_transport, + ) + with pytest.raises(TypeError, match="Unsupported stream type"): + session.upload(stream=invalid_stream) + + +def test_upload_state_properties(): + state = upload_state._ProtocolState("https://api.example.com/init", chunk_size=500) + assert state.initial_url == "https://api.example.com/init" + assert state.upload_url == "https://api.example.com/init" + assert state.bytes_uploaded == 0 + assert state.total_bytes is None + assert state.finished is False + assert state.invalid is False + assert state.chunk_size == 500 + + # With granularity alignment + state._chunk_granularity = 256 + assert state.chunk_size == 512 + + +def test_upload_state_start_errors(): + empty_state = upload_state._ProtocolState() + with pytest.raises( + ValueError, match="upload_url must be provided to start an upload" + ): + empty_state.build_start_request() + + state = upload_state._ProtocolState("https://api.example.com/init") + with pytest.raises(ValueError, match="Start command failed with status 500"): + state.process_start_response(500, {}) + assert state.invalid is True + + state2 = upload_state._ProtocolState("https://api.example.com/init") + with pytest.raises(ValueError, match="Server did not return"): + state2.process_start_response(200, {"X-Goog-Upload-Status": "active"}) + assert state2.invalid is True + + +def test_upload_state_chunk_and_query_errors(): + state = upload_state._ProtocolState() + with pytest.raises(ValueError, match="Upload session URL not established"): + state.build_chunk_request(b"data", is_last_chunk=True) + + with pytest.raises(ValueError, match="Upload session URL not established"): + state.build_query_request() + + with pytest.raises(ValueError, match="Upload session URL not established"): + state.build_cancel_request() + + # process_chunk_response with non-200/201 status code + state.process_chunk_response(503, {}, 10) + assert state.bytes_uploaded == 0 + + # process_query_response with non-200/201 status code + with pytest.raises(ValueError, match="Query recovery failed with status 500"): + state.process_query_response(500, {}) + assert state.invalid is True + + # process_query_response with final status + state3 = upload_state._ProtocolState("https://api.example.com/init") + state3.process_query_response(200, {"X-Goog-Upload-Status": "final"}) + assert state3.finished is True + + # process_query_response with cancelled status + state4 = upload_state._ProtocolState("https://api.example.com/init") + with pytest.raises(UploadCancelledError): + state4.process_query_response(200, {"X-Goog-Upload-Status": "cancelled"}) + assert state4.invalid is True + + +def test_sync_upload_session_properties_and_enrichment(): + session_transport = mock.create_autospec(requests.Session, instance=True) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + transport=session_transport, + ) + assert session._get_transport(None) is session_transport + assert session._state.upload_url == "https://api.example.com/init" + assert session.bytes_uploaded == 0 + assert session._state.total_bytes is None + assert session.finished is False + assert session._state.invalid is False + + # Exception without __dict__ does not fail _enrich_exception + exc_no_dict = Exception() + session._enrich_exception(exc_no_dict) + + +def test_sync_upload_session_transport_missing(): + session = ResumableUploadSession(upload_url="https://api.example.com/init") + with pytest.raises( + ValueError, match="A requests.Session transport must be provided" + ): + session.upload(stream=b"payload") + + with pytest.raises( + ValueError, match="A requests.Session transport must be provided" + ): + session.cancel() + + +def test_sync_deadline_handling(): + past = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=10) + config = ResumableUploadConfig(deadline=past) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config, + ) + with pytest.raises(exceptions.DeadlineExceeded): + session._get_deadline_remaining() + + future_naive = datetime.datetime.now() + datetime.timedelta(hours=1) + config2 = ResumableUploadConfig(deadline=future_naive) + session2 = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config2, + ) + rem = session2._get_deadline_remaining() + assert rem is not None and rem > 0 + + +def test_sync_retry_predicate_branches(): + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + ) + pred_transfer = session._get_retry_predicate(is_start=False) + pred_start = session._get_retry_predicate(is_start=True) + + assert pred_transfer(exceptions.DeadlineExceeded("deadline")) is False + assert pred_transfer(TransferStalledError("stalled")) is False + assert pred_transfer(UploadCancelledError("cancelled")) is False + assert pred_transfer(MissingStatusHeaderError("missing")) is True + assert pred_start(MissingStatusHeaderError("missing")) is False + assert pred_transfer(requests.exceptions.ConnectionError("conn")) is True + assert pred_transfer(requests.exceptions.ChunkedEncodingError("chunked")) is True + assert pred_transfer(exceptions.from_http_status(503, "503")) is True + assert pred_transfer(exceptions.from_http_status(400, "400")) is True + assert pred_transfer(exceptions.from_http_status(412, "412")) is True + assert pred_transfer(exceptions.from_http_status(416, "416")) is True + assert pred_start(exceptions.from_http_status(400, "400")) is False + assert pred_start(exceptions.from_http_status(412, "412")) is False + assert pred_start(exceptions.from_http_status(416, "416")) is False + assert pred_transfer(exceptions.from_http_status(409, "409")) is False + assert pred_start(exceptions.from_http_status(409, "409")) is False + assert pred_transfer(exceptions.from_http_status(403, "403")) is False + assert pred_transfer(TypeError("other")) is False + + +def test_sync_reposition_stream_errors(): + session_transport = mock.create_autospec(requests.Session, instance=True) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + transport=session_transport, + ) + + unseekable = mock.Mock() + unseekable.seekable.return_value = False + with pytest.raises(UnseekableStreamError, match="Stream is not seekable"): + session._reposition_stream_offset(unseekable, 100) + + failing_seek = mock.Mock() + failing_seek.seekable.return_value = True + failing_seek.seek.side_effect = OSError("Disk read failure") + with pytest.raises(UnseekableStreamError, match="Failed to seek stream"): + session._reposition_stream_offset(failing_seek, 100) + + +def test_sync_prepare_stream_seekable_and_iterable(): + session_transport = mock.create_autospec(requests.Session, instance=True) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + transport=session_transport, + ) + + stream_obj, computed_size = session._prepare_stream([b"hello ", b"world"], None) + assert stream_obj.seekable() is False + assert stream_obj.read(4) == b"hell" + assert stream_obj.read(4) == b"o wo" + assert stream_obj.read(4) == b"rld" + assert stream_obj.read(4) == b"" + assert computed_size is None + + stream_obj_all, _ = session._prepare_stream([b"hello ", b"world"], None) + assert stream_obj_all.read() == b"hello world" + + class CustomSeekable: + def __init__(self, data: bytes): + self._bio = io.BytesIO(data) + + def read(self, n: int = -1) -> bytes: + return self._bio.read(n) + + def seek(self, offset: int, whence: int = io.SEEK_SET) -> int: + return self._bio.seek(offset, whence) + + def tell(self) -> int: + return self._bio.tell() + + def seekable(self) -> bool: + return True + + custom = CustomSeekable(b"0123456789") + custom.seek(2) + stream_obj2, computed_size2 = session._prepare_stream(custom, None) + assert computed_size2 == 8 + assert custom.tell() == 2 + + +def test_sync_config_headers(): + cfg_dict = ResumableUploadConfig(headers={"X-Key": "Val"}) + assert cfg_dict.start_headers == [("X-Key", "Val")] + + cfg_list = ResumableUploadConfig(headers=[("X-Key", "Val")]) + assert cfg_list.start_headers == [("X-Key", "Val")] + + cfg_none = ResumableUploadConfig(headers=None) + assert cfg_none.start_headers is None + + +def test_sync_cancel_failure_raises(): + session_transport = mock.create_autospec(requests.Session, instance=True) + err_resp = mock.create_autospec(requests.Response, instance=True) + err_resp.ok = False + err_resp.status_code = 500 + err_resp.headers = {} + err_resp.request = mock.Mock(method="POST", url="https://upload.example.com") + err_resp.json.return_value = {"error": {"message": "Server Error", "errors": []}} + session_transport.request.return_value = err_resp + + session = ResumableUploadSession( + upload_url="https://upload.example.com/resumable-123", + transport=session_transport, + ) + with pytest.raises(exceptions.GoogleAPICallError): + session.cancel() + + +def test_sync_resume_chunk_size_override(): + session_transport = mock.create_autospec(requests.Session, instance=True) + query_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "0", + }, + ) + chunk_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b"done", + ) + session_transport.request.side_effect = [query_resp, chunk_resp] + + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + transport=session_transport, + ) + session.resume( + upload_url="https://upload.example.com/resumable-123", + stream=b"data", + chunk_size=1024, + ) + assert session.chunk_size == 1024 + + +def test_sync_on_progress_and_capture(): + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + ) + session._state._upload_url = "https://api.example.com/init" + progress_queue = [] + session._notify_progress( + common.ProgressState.UPLOADING, progress_queue=progress_queue + ) + assert len(progress_queue) == 1 + assert progress_queue[0].state == common.ProgressState.UPLOADING + + +def test_sync_naive_deadline_tz(): + naive = datetime.datetime.now() + datetime.timedelta(hours=1) + config = ResumableUploadConfig(deadline=naive) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config, + ) + rem = session._get_deadline_remaining() + assert rem is not None and rem > 0 + assert session._get_start_timeout() <= rem + + +def test_sync_get_retry_start_and_override(): + """Verifies how user-supplied Retry objects are resolved and wrapped.""" + + class CustomApiError(Exception): + """Example API-specific transient exception provided by a caller.""" + + # ------------------------------------------------------------------------- + # Scenario 1: User provides a custom predicate to retry an API-specific error + # ------------------------------------------------------------------------- + custom_start_retry = google.api_core.retry.Retry( + initial=0.5, + predicate=lambda exc: isinstance(exc, CustomApiError), + ) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + start_retry=custom_start_retry, + ) + resolved_start = session._get_retry() + + # Backoff settings (0.5s initial delay) are preserved + assert resolved_start._initial == 0.5 + # The user's custom exception is retried, while unrelated errors are not + assert resolved_start._predicate(CustomApiError("rate limit")) is True + assert resolved_start._predicate(ValueError("invalid input")) is False + + # ------------------------------------------------------------------------- + # Scenario 2: Terminal errors are blocked even if custom predicate returns True + # ------------------------------------------------------------------------- + overly_broad_retry = google.api_core.retry.Retry( + initial=0.25, + predicate=lambda exc: True, + ) + resolved_override = session._get_retry(retry_override=overly_broad_retry) + + assert resolved_override._initial == 0.25 + # Terminal transfer errors must never be retried, preventing infinite loops + assert ( + resolved_override._predicate(exceptions.DeadlineExceeded("deadline expired")) + is False + ) + assert ( + resolved_override._predicate(exceptions.TransferStalledError("upload stalled")) + is False + ) + + # ------------------------------------------------------------------------- + # Scenario 3: Category 2 recovery (400, 412, 416, MissingStatusHeaderError) + # is preserved during chunk transfer + # ------------------------------------------------------------------------- + default_stream = session._get_streaming_retry() + assert default_stream.timeout is None + + restrictive_chunk_retry = google.api_core.retry.StreamingRetry( + initial=0.1, + predicate=lambda exc: False, + ) + resolved_chunk = session._get_streaming_retry( + retry_override=restrictive_chunk_retry + ) + + assert resolved_chunk._initial == 0.1 + # Category 2 recovery triggers (400, 412, 416, and missing status header) + # must still return True during chunk transfer so the client can query + # server offset and synchronize state + assert ( + resolved_chunk._predicate(exceptions.from_http_status(400, "Bad Request")) + is True + ) + assert ( + resolved_chunk._predicate( + exceptions.from_http_status(412, "Precondition Failed") + ) + is True + ) + assert ( + resolved_chunk._predicate( + exceptions.from_http_status(416, "Range Not Satisfiable") + ) + is True + ) + missing_header_error = MissingStatusHeaderError("Missing X-Goog-Upload-Status") + assert resolved_chunk._predicate(missing_header_error) is True + # Unretriable HTTP status codes (such as 409 Conflict) and non-protocol errors + # follow the predicate and return False + assert ( + resolved_chunk._predicate(exceptions.from_http_status(409, "Conflict")) is False + ) + assert resolved_chunk._predicate(RuntimeError("unexpected crash")) is False + + +def test_sync_stall_control_with_deadline(): + config = ResumableUploadConfig( + stall_minimum_rate=1024, + stall_timeout=10.0, + ) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config, + ) + # 512 bytes / 1024 B/s -> expected_sec = 0.5s. + # With aggregate_lag = 0.0 and stall_timeout = 10.0s: + # next_chunk_timeout = 0.5 - 0.0 + 10.0 = 10.5s. + # timeout_override=15.0 exceeds next_chunk_timeout, so it clamps to 10.5s. + t1 = session._compute_chunk_timeout(512, timeout_override=15.0) + assert t1 == pytest.approx(10.5) + # timeout_override=8.0 is below next_chunk_timeout (10.5s), so 8.0s is used + # directly instead of the default heuristic (max(5.0, 2 * 0.5) = 5.0s). + assert session._compute_chunk_timeout(512, timeout_override=8.0) == pytest.approx( + 8.0 + ) + + # When stall control is disabled, timeout_override is returned as-is. + session_no_stall = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(stall_minimum_rate=0, stall_timeout=0), + ) + assert session_no_stall._compute_chunk_timeout( + 512, timeout_override=15.0 + ) == pytest.approx(15.0) + + # Overall upload deadline (5s left) caps the per-attempt timeout. + config2 = ResumableUploadConfig( + stall_minimum_rate=1024, + stall_timeout=10.0, + deadline=datetime.datetime.now(datetime.timezone.utc) + + datetime.timedelta(seconds=5), + ) + session2 = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config2, + ) + t2 = session2._compute_chunk_timeout(512) + assert t2 <= 5.0 + + # If the upload stalls and the overall deadline has already passed, + # DeadlineExceeded takes precedence over TransferStalledError. + config3 = ResumableUploadConfig( + stall_minimum_rate=1024, + stall_timeout=10.0, + deadline=datetime.datetime.now(datetime.timezone.utc) + - datetime.timedelta(seconds=5), + ) + session3 = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config3, + ) + with pytest.raises(exceptions.DeadlineExceeded): + session3._update_stall_control(512, 15.0) + + # 512 bytes / 1024 B/s -> expected_sec = 0.5s. Taking 15.0s yields + # current_lag = 14.5s, which exceeds stall_timeout (10.0s). + config4 = ResumableUploadConfig( + stall_minimum_rate=1024, + stall_timeout=10.0, + ) + session4 = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config4, + ) + session4._state._upload_url = "https://api.example.com/init" + with pytest.raises(exceptions.TransferStalledError): + session4._update_stall_control(512, 15.0) + + # 10240 bytes / 1024 B/s -> expected_sec = 10.0s. + # Taking 10.5s gives current_lag = 10.5 - 10.0 = 0.5s. + # Only the 0.5s lag counts toward stall_timeout (10.0s), not the full 10.5s, + # so the transfer does not stall yet. + session5 = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config4, + ) + session5._state._upload_url = "https://api.example.com/init" + session5._update_stall_control(10240, 10.5) + assert session5._aggregate_lag == pytest.approx(0.5) + assert session5._stall_timeout_started is not None + session5._buffered_chunk = memoryview(b"stale") + session5._reset_transfer_state() + assert session5._buffered_chunk is None + assert session5._aggregate_lag == 0.0 + assert session5._stall_timeout_started is None + + +def test_sync_update_stall_control_disabled(): + config = ResumableUploadConfig(stall_minimum_rate=0, stall_timeout=10.0) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config, + ) + session._update_stall_control(512, 5.0) + assert session._aggregate_lag == 0.0 + + # Ensure that _stall_timeout_started resets to None when transfer rate exceeds minimum rate (no lag). + active_config = ResumableUploadConfig(stall_minimum_rate=1024, stall_timeout=10.0) + active_session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=active_config, + ) + active_session._stall_timeout_started = 100.0 + active_session._update_stall_control(1024, 0.1) + assert active_session._aggregate_lag == 0.0 + assert active_session._stall_timeout_started is None + + +def test_sync_initiate_failure(): + transport = mock.create_autospec(requests.Session, instance=True) + resp = mock.create_autospec(requests.Response, instance=True) + resp.ok = False + resp.status_code = 400 + resp.headers = {} + resp.json.return_value = {"error": {"message": "Init Failed"}} + resp.request = mock.Mock(method="POST", url="https://api.example.com/init") + transport.request.return_value = resp + + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + transport=transport, + ) + with pytest.raises(exceptions.GoogleAPICallError): + session._initiate(transport=transport) + + +def test_sync_transmit_empty_stream(): + transport = mock.create_autospec(requests.Session, instance=True) + resp = mock.create_autospec(requests.Response, instance=True) + resp.ok = True + resp.status_code = 200 + resp.headers = {"X-Goog-Upload-Status": "final"} + resp.content = b"done" + transport.request.return_value = resp + + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + transport=transport, + ) + stream = io.BytesIO(b"") + session._state._upload_url = "https://upload.example.com/resumable-123" + result = session._transmit_chunk(transport, stream, size=0) + assert result is resp + + +def test_sync_transmit_chunk_timeout_with_stall_control_active(monkeypatch): + transport = mock.create_autospec(requests.Session, instance=True) + transport.request.side_effect = requests.exceptions.Timeout("Read timeout") + + config = ResumableUploadConfig( + stall_minimum_rate=1024, + stall_timeout=10.0, + ) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config, + transport=transport, + ) + session._state._upload_url = "https://upload.example.com/resumable-123" + + # Attempt 1 times out at 5.0s (< stall_timeout=10.0s): re-raises Timeout so retry/recovery can run + clock_vals = iter([0.0, 5.0, 5.0, 5.0, 10.0, 10.0, 10.0, 10.0]) + monkeypatch.setattr(upload, "_monotonic_clock", lambda: next(clock_vals)) + with pytest.raises(requests.exceptions.Timeout): + session._transmit_chunk(transport, io.BytesIO(b"data"), size=4) + + # Attempt 2 reaches elapsed stall duration of 10.0s (>= stall_timeout=10.0s): raises TransferStalledError + with pytest.raises(exceptions.TransferStalledError): + session._transmit_chunk(transport, io.BytesIO(b"data"), size=4) + + config_dl = ResumableUploadConfig( + stall_minimum_rate=1024, + stall_timeout=10.0, + deadline=datetime.datetime.now(datetime.timezone.utc) + + datetime.timedelta(seconds=5), + ) + session_dl = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config_dl, + transport=transport, + ) + session_dl._state._upload_url = "https://upload.example.com/resumable-123" + session_dl._get_deadline_remaining = mock.Mock( + side_effect=[5.0, exceptions.DeadlineExceeded("Deadline exceeded")] + ) + with pytest.raises(exceptions.DeadlineExceeded): + session_dl._transmit_chunk(transport, io.BytesIO(b"data"), size=4) + + +def test_sync_per_attempt_timeout_retries_before_stall_timeout(monkeypatch): + """Verifies that hitting per_attempt_timeout (5s) retries via recovery rather + than prematurely raising TransferStalledError when stall_timeout (120s) has not elapsed.""" + transport = mock.create_autospec(requests.Session, instance=True) + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-123", + }, + ) + query_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "0", + }, + ) + final_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b'{"status": "completed"}', + ) + transport.request.side_effect = [ + start_resp, + requests.exceptions.Timeout("Per-attempt 5s timeout"), + query_resp, + final_resp, + ] + + # Configure stall control with 120s stall_timeout; per_attempt_timeout is 5.0s + config = ResumableUploadConfig( + stall_minimum_rate=1024, + stall_timeout=120.0, + ) + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=transport, + ) + + # First attempt times out after 5.0s (< 120.0s stall_timeout); retry succeeds at 6.0s + clock_vals = iter([0.0, 5.0, 5.0, 5.0, 6.0, 6.0]) + monkeypatch.setattr(upload, "_monotonic_clock", lambda: next(clock_vals)) + + result = session.upload( + stream=b"test data", + retry=google.api_core.retry.StreamingRetry(initial=0.01, maximum=0.01), + ) + assert result == b'{"status": "completed"}' + assert transport.request.call_count == 4 + + +def test_sync_transmit_chunk_timeout_outer_exception(): + transport = mock.create_autospec(requests.Session, instance=True) + transport.request.side_effect = requests.exceptions.Timeout("Read timeout") + + config = ResumableUploadConfig( + stall_minimum_rate=0, + ) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config, + transport=transport, + ) + session._state._upload_url = "https://upload.example.com/resumable-123" + no_retry = google.api_core.retry.StreamingRetry( + predicate=lambda e: False, timeout=0 + ) + with pytest.raises(exceptions.TransferStalledError): + list( + session._transmit_all_chunks( + transport, io.BytesIO(b"data"), 4, retry=no_retry + ) + ) + + # To hit outer exception handler in _transmit_all_chunks with elapsed deadline + config_dl = ResumableUploadConfig( + stall_minimum_rate=0, + deadline=datetime.datetime.now(datetime.timezone.utc) + + datetime.timedelta(seconds=5), + ) + session_dl = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config_dl, + transport=transport, + ) + session_dl._state._upload_url = "https://upload.example.com/resumable-123" + session_dl._get_deadline_remaining = mock.Mock( + side_effect=[5.0, 5.0, exceptions.DeadlineExceeded("Deadline exceeded")] + ) + with pytest.raises(exceptions.DeadlineExceeded): + list( + session_dl._transmit_all_chunks( + transport, io.BytesIO(b"data"), 4, retry=no_retry + ) + ) + + # Non-timeout RetryError re-raises RetryError rather than TransferStalledError + transport_conn = mock.create_autospec(requests.Session, instance=True) + transport_conn.request.side_effect = requests.exceptions.ConnectionError( + "Connection dropped" + ) + session_conn = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config, + transport=transport_conn, + ) + session_conn._state._upload_url = "https://upload.example.com/resumable-123" + with pytest.raises(exceptions.RetryError): + list( + session_conn._transmit_all_chunks( + transport_conn, io.BytesIO(b"data"), 4, retry=no_retry + ) + ) + + +def test_sync_recover_failure(): + transport = mock.create_autospec(requests.Session, instance=True) + resp = mock.create_autospec(requests.Response, instance=True) + resp.ok = False + resp.status_code = 400 + resp.headers = {} + resp.json.return_value = {"error": {"message": "Recovery Failed"}} + resp.request = mock.Mock( + method="POST", url="https://upload.example.com/resumable-123" + ) + transport.request.return_value = resp + + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + transport=transport, + ) + session._state._upload_url = "https://upload.example.com/resumable-123" + with pytest.raises(exceptions.GoogleAPICallError): + session._recover(transport, io.BytesIO(b"data")) + + +def test_sync_transmit_all_chunks_completed_without_response(): + transport = mock.create_autospec(requests.Session, instance=True) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + transport=transport, + ) + session._state._finished = True + with pytest.raises( + ValueError, match="Upload completed without receiving a final response" + ): + list(session._transmit_all_chunks(transport, io.BytesIO(b"data"), 4)) + + +def test_sync_iter_resume_errors(): + transport = mock.create_autospec(requests.Session, instance=True) + session = ResumableUploadSession(transport=transport) + with pytest.raises(ValueError, match="An upload URL must be provided to resume"): + list(session.iter_resume(upload_url=None, stream=b"data")) + + with pytest.raises( + ValueError, match="A data stream or payload must be provided to resume" + ): + list( + session.iter_resume(upload_url="https://api.example.com/init", stream=None) + ) + + +def test_sync_prepare_stream_tell_error(): + class TellFailingStream(io.BytesIO): + def tell(self) -> int: + raise OSError("Tell failed") + + session = ResumableUploadSession() + stream = TellFailingStream(b"data") + stream_obj, computed_size = session._prepare_stream(stream, None) + assert session._start_stream_offset == 0 + + +def test_sync_format_response_payload_custom_inputs(): + from google.api_core.resumable_transfer.upload import _format_response_payload + + class CustomBytesConvertible: + def __bytes__(self) -> bytes: + return b"custom_bytes" + + res = _format_response_payload(CustomBytesConvertible(), response_type=None) + assert res == b"custom_bytes" + + res_parsed = _format_response_payload( + CustomBytesConvertible(), response_type=lambda x: x + b"_extra" + ) + assert res_parsed == b"custom_bytes_extra" + + from google.protobuf import empty_pb2 + + msg_instance = empty_pb2.Empty() + res_msg = _format_response_payload(b"{}", response_type=msg_instance) + assert isinstance(res_msg, empty_pb2.Empty) + + +class SlotException(Exception): + __slots__ = () + + +class NoTellStream: + def read(self, n): + return b"" + + +class TellErrorStream: + def read(self, n): + return b"" + + def tell(self): + raise OSError("tell failed") + + +class CustomReadStream: + def __init__(self, data): + self.data = data + + def read(self, n): + return self.data + + +def test_sync_enrich_exception(): + session = ResumableUploadSession( + upload_url="https://upload.example.com/resumable-123", + transport=mock.sentinel.transport, + ) + exc = RuntimeError("test error") + session._enrich_exception(exc) + assert getattr(exc, "upload_url") == "https://upload.example.com/resumable-123" + + +def test_sync_notify_progress_no_upload_url(): + session = ResumableUploadSession( + upload_url=None, + transport=mock.sentinel.transport, + ) + session._notify_progress(common.ProgressState.STARTED) + + +def test_sync_should_retry_request_exception_not_retryable(): + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + transport=mock.sentinel.transport, + ) + should_retry = session._get_retry_predicate() + exc = requests.exceptions.HTTPError("Non-retryable HTTP Error") + assert not should_retry(exc) + + +class NoSeekableButSeekStream: + def read(self, n): + return b"" + + def seek(self, offset): + pass + + +class UnseekableReadStream: + def read(self, n): + return b"" + + def seekable(self): + return False + + +class SeekableNoTellStream: + def read(self, n): + return b"" + + def seekable(self): + return True + + +def test_sync_rewind_stream_no_seekable_attr(): + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + transport=mock.sentinel.transport, + ) + stream = NoSeekableButSeekStream() + session._reposition_stream_offset(stream, 0) + + +def test_sync_recover_buffered_chunk_out_of_bounds(): + session = ResumableUploadSession() + session._buffered_chunk = memoryview(b"data") + session._buffered_chunk_offset = 0 + + unseekable = mock.Mock() + unseekable.seekable.return_value = False + + with pytest.raises(UnseekableStreamError): + session._reposition_stream_offset(unseekable, 10) + + assert session._buffered_chunk is None + + +def test_sync_prepare_stream_edge_cases(): + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + transport=mock.sentinel.transport, + ) + + # 1. NoTellStream: Has read, but no tell + stream1 = NoTellStream() + s_obj1, size1 = session._prepare_stream(stream1, size=None) + assert s_obj1 is stream1 + assert size1 is None + assert session._start_stream_offset == 0 + + # 2. TellErrorStream: Has read and tell, but tell raises OSError + stream2 = TellErrorStream() + s_obj2, size2 = session._prepare_stream(stream2, size=None) + assert s_obj2 is stream2 + assert size2 is None + assert session._start_stream_offset == 0 + + # 3. CustomReadStream: Has read, no getbuffer, no seekable + stream3 = CustomReadStream(b"hello") + s_obj3, size3 = session._prepare_stream(stream3, size=None) + assert s_obj3 is stream3 + assert size3 is None + + # 4. UnseekableReadStream: Has read, has seekable returning False + stream4 = UnseekableReadStream() + s_obj4, size4 = session._prepare_stream(stream4, size=None) + assert s_obj4 is stream4 + assert size4 is None + + # 5. SeekableNoTellStream: Has read, has seekable returning True, no tell + stream5 = SeekableNoTellStream() + s_obj5, size5 = session._prepare_stream(stream5, size=None) + assert s_obj5 is stream5 + assert size5 is None + + +def test_sync_transmit_all_chunks_captured_empty(): + session_transport = mock.create_autospec(requests.Session, instance=True) + chunk_resp = mock.create_autospec(requests.Response, instance=True) + chunk_resp.status_code = 200 + chunk_resp.content = b"{}" + chunk_resp.headers = {"X-Goog-Upload-Status": "final"} + session_transport.request.return_value = chunk_resp + + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + ) + session._state._upload_url = "https://upload.example.com/resumable-123" + + stream_obj = io.BytesIO(b"data") + + # Consume the generator + list( + session._transmit_all_chunks( + session_transport, stream_obj, 4, progress_queue=None + ) + ) + + +def test_sync_upload_multiple_chunks(): + session_transport = mock.create_autospec(requests.Session, instance=True) + + # Start response + start_resp = mock.create_autospec(requests.Response, instance=True) + start_resp.status_code = 200 + start_resp.content = b"" + start_resp.headers = { + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-123", + } + + # 1st chunk response: active + chunk1_resp = mock.create_autospec(requests.Response, instance=True) + chunk1_resp.status_code = 200 + chunk1_resp.content = b"" + chunk1_resp.headers = {"X-Goog-Upload-Status": "active"} + + # 2nd chunk response: final + chunk2_resp = mock.create_autospec(requests.Response, instance=True) + chunk2_resp.status_code = 200 + chunk2_resp.content = b"{}" + chunk2_resp.headers = {"X-Goog-Upload-Status": "final"} + + session_transport.request.side_effect = [start_resp, chunk1_resp, chunk2_resp] + + config = ResumableUploadConfig(chunk_size=5) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config, + ) + + res = session.upload(stream=b"0123456789", transport=session_transport) + assert res == chunk2_resp.content + assert session.bytes_uploaded == 10 + assert session_transport.request.call_count == 3 + + +def test_sync_format_response_payload_unsupported_type(): + from google.api_core.resumable_transfer.upload import _format_response_payload + + res = _format_response_payload(b"some content", "unsupported") + assert res == b"some content" + + +def test_sync_prepare_stream_explicit_size(): + session = ResumableUploadSession(upload_url="https://api.example.com/init") + + # 1. bytes with explicit size + _, computed_size = session._prepare_stream(b"abcd", size=4) + assert computed_size == 4 + + # 2. Iterable with explicit size + _, computed_size = session._prepare_stream([b"ab", b"cd"], size=4) + assert computed_size == 4 + + # 3. BinaryIO with explicit size + _, computed_size = session._prepare_stream(io.BytesIO(b"abcd"), size=4) + assert computed_size == 4 + + +def test_state_process_chunk_response_unknown_status(): + from google.api_core.resumable_transfer.upload_state import _ProtocolState + + state = _ProtocolState(upload_url="https://api.example.com/init") + state.process_chunk_response(200, {"X-Goog-Upload-Status": "unknown"}, 100) + assert state.bytes_uploaded == 0 + assert not state.finished + + +def test_sync_partial_chunk_recovery_does_not_prematurely_finalize(): + """Ensure that retrying a partially committed chunk (len < chunk_size) does not prematurely finalize. + + When a 4-byte chunk (b"0123") partially succeeds (server commits 2 bytes) + and is retried, ensure that the remaining 2 bytes (b"23") are sent with + "upload" rather than "upload, finalize" so the remaining payload (b"45") + is not dropped. + """ + server_received_bytes = bytearray() + call_count = 0 + + def handle_request(method, url, data=None, headers=None, **kwargs): + nonlocal call_count + call_count += 1 + cmd = headers.get("X-Goog-Upload-Command", "") if headers else "" + + resp = mock.create_autospec(requests.Response, instance=True) + resp.request = mock.Mock(method=method, url=url) + resp.content = b"" + + # Request 1: start session + if call_count == 1: + assert cmd == "start" + resp.status_code = 200 + resp.ok = True + resp.headers = { + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-123", + } + return resp + + # Request 2: initial Chunk 1 (b"0123") -> server commits partial 2 bytes (b"01"), then fails 503 + if call_count == 2: + assert cmd == "upload" + assert bytes(data) == b"0123" + server_received_bytes.extend(b"01") + resp.status_code = 503 + resp.ok = False + resp.headers = {} + resp.json.return_value = { + "error": {"code": 503, "message": "Service Unavailable"} + } + return resp + + # Request 3: recovery query -> server reports 2 committed bytes + if call_count == 3: + assert cmd == "query" + resp.status_code = 200 + resp.ok = True + resp.headers = { + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": str(len(server_received_bytes)), + } + return resp + + # Subsequent upload requests (Request 4: remaining b"23", Request 5: final b"45") + if data: + server_received_bytes.extend(bytes(data)) + + resp.status_code = 200 + resp.ok = True + if "finalize" in cmd: + resp.headers = {"X-Goog-Upload-Status": "final"} + resp.content = b"{}" + else: + resp.headers = {"X-Goog-Upload-Status": "active"} + return resp + + session_transport = mock.create_autospec(requests.Session, instance=True) + session_transport.request.side_effect = handle_request + + # Ensure that the unified outer StreamingRetry coordinates backoff and + # triggers protocol-level _recover() on retryable errors before retransmitting. + retry_cfg = google.api_core.retry.StreamingRetry( + predicate=ResumableUploadSession()._get_retry_predicate(), + initial=0.001, + ) + config = ResumableUploadConfig(chunk_size=4) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config, + ) + + session.upload(stream=b"012345", transport=session_transport, retry=retry_cfg) + + # Verify no data loss occurred: server must receive all 6 bytes (b"012345"), not truncated b"0123" + assert bytes(server_received_bytes) == b"012345" + + +def test_sync_streaming_retry_and_recovery_final(): + """Ensure that StreamingRetry configuration and recovery returning 'final' status are handled properly.""" + session_transport = mock.create_autospec(requests.Session, instance=True) + + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/resumable-123", + }, + ) + err503_resp = mock.Mock( + ok=False, + status_code=503, + headers={}, + ) + err503_resp.json.return_value = {"error": {"code": 503, "message": "Unavailable"}} + err503_resp.text = '{"error": {"code": 503, "message": "Unavailable"}}' + err503_resp.content = err503_resp.text.encode("utf-8") + + # Recovery query reports that the upload already reached 'final' status on the server + query_final_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b'{"name": "already_done.txt", "size": 4}', + ) + + session_transport.request.side_effect = [ + start_resp, + err503_resp, + query_final_resp, + ] + + streaming_retry = google.api_core.retry.StreamingRetry( + predicate=ResumableUploadSession()._get_retry_predicate(), + initial=0.001, + ) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + response_type=DummyResponse, + ) + resolved_streaming = session._get_streaming_retry(retry_override=streaming_retry) + assert resolved_streaming._initial == 0.001 + + result = session.upload( + stream=b"data", transport=session_transport, retry=streaming_retry + ) + assert isinstance(result, DummyResponse) + assert result.name == "already_done.txt" + assert session.finished is True + + +def test_sync_method_override_arguments() -> None: + """Verifies content_type and on_progress overrides on initiate, upload, and resume.""" + # 1. initiate with content_type override + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/123", + }, + content=b"", + ) + transport1 = mock.Mock(spec=requests.Session) + transport1.request.return_value = start_resp + session1 = ResumableUploadSession(upload_url="https://api.example.com/start") + session1._initiate(transport=transport1, content_type="text/plain") + assert session1._content_type == "text/plain" + + # 2. upload with content_type override + chunk_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b"{}", + ) + transport2 = mock.Mock(spec=requests.Session) + transport2.request.side_effect = [start_resp, chunk_resp] + session2 = ResumableUploadSession(upload_url="https://api.example.com/start") + session2.upload( + stream=b"data", + transport=transport2, + content_type="text/csv", + ) + assert session2._content_type == "text/csv" + + # 3. resume with chunk_size override + query_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "0", + }, + content=b"", + ) + transport3 = mock.Mock(spec=requests.Session) + transport3.request.side_effect = [query_resp, chunk_resp] + session3 = ResumableUploadSession() + session3.resume( + upload_url="https://upload.example.com/123", + stream=b"data", + transport=transport3, + ) diff --git a/packages/google-api-core/tests/unit/test_resumable_transfer_scenarios.py b/packages/google-api-core/tests/unit/test_resumable_transfer_scenarios.py new file mode 100644 index 000000000000..d7729a01abe6 --- /dev/null +++ b/packages/google-api-core/tests/unit/test_resumable_transfer_scenarios.py @@ -0,0 +1,862 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Test suite covering edge cases, stream boundaries, retry interactions, and sync/async parity. + +These tests evaluate: +1. Stream continuity and subsequent chunk delivery following full/partial chunk recovery. +2. Parity in return types, retry configuration, and exception handling between sync and async. +3. Memory efficiency when consuming streaming iterables. +4. Coordination of retry policies, backoffs, and timeout handling during transport interruptions. +""" + +import asyncio +import datetime +from unittest import mock + +import pytest +import requests + +import google.api_core.retry +from google.api_core.resumable_transfer.upload import ( + ResumableUploadConfig, + ResumableUploadSession, +) +from google.api_core.resumable_transfer.upload_async import ( + AsyncResumableUploadSession, +) + +try: + import aiohttp # noqa: F401 + import google.auth.aio.transport # noqa: F401 + + GOOGLE_AUTH_AIO_INSTALLED = True +except ImportError: + GOOGLE_AUTH_AIO_INSTALLED = False + + +class _MockAiohttpResp: + """Mock for aiohttp response context manager.""" + + def __init__(self, status: int, headers: dict, body: bytes): + self.status = status + self.headers = headers + self._body = body + + async def read(self) -> bytes: + return self._body + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + +# ============================================================================== +# 1. Stream Recovery & Chunk Continuity +# ============================================================================== + + +def test_full_chunk_recovery_transmits_subsequent_chunks_sync(): + """Verifies that sync upload continues transmitting subsequent chunks after full chunk recovery. + + Why this test is needed: + When chunk 1 (bytes 0..100) fails with a recoverable error (e.g. 412 Precondition Failed) + after the server has actually committed all 100 bytes, the status query returns + ``X-Goog-Upload-Size-Received: 100``. Slicing the active in-memory buffer by the + committed byte count (`_buffered_chunk[100:]`) yields a 0-length ``memoryview``. + If `_reposition_stream_offset` leaves a non-``None`` 0-length buffer in place instead + of resetting `_buffered_chunk = None`, the next transmission attempt treats that empty + buffer as the active chunk rather than reading chunk 2 (bytes 100..200) from the stream. + """ + transport = mock.Mock() + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/session1", + }, + ) + resp_412 = mock.Mock( + ok=False, + status_code=412, + text="Precondition Failed", + content=b"Precondition Failed", + headers={}, + json=lambda: {}, + ) + # Status query confirms the entire first chunk (100 bytes) was committed on the server. + query_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "100", + }, + ) + final_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b"done", + ) + transport.request.side_effect = [start_resp, resp_412, query_resp, final_resp] + + # Total payload: 200 bytes across two 100-byte chunks. + payload = b"A" * 100 + b"B" * 100 + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(chunk_size=100), + transport=transport, + ) + session.upload(stream=payload) + + transmitted_chunks = [ + request_call.kwargs.get("data") + for request_call in transport.request.call_args_list + if request_call.kwargs.get("headers", {}).get("X-Goog-Upload-Command") + in ("upload", "upload, finalize") + ] + total_bytes_sent = sum( + len(chunk) for chunk in transmitted_chunks if chunk is not None + ) + assert total_bytes_sent == 200, ( + f"Expected 200 bytes uploaded across all chunks, but received {total_bytes_sent} " + f"bytes. Transmitted chunks: {transmitted_chunks}" + ) + + +@pytest.mark.skipif( + not GOOGLE_AUTH_AIO_INSTALLED, + reason="Skipped because google-api-core[async_rest] is not installed", +) +@pytest.mark.asyncio +async def test_full_chunk_recovery_transmits_subsequent_chunks_async(): + """Async counterpart verifying subsequent chunks continue transmitting after full chunk recovery. + + Why this test is needed: + Ensures ``AsyncResumableUploadSession._recover`` clears ``_buffered_chunk = None`` + when a status query confirms all bytes of the active chunk were committed, so the + next iteration reads the next chunk from the stream instead of sending an empty slice. + """ + transport = mock.Mock() + start_resp = _MockAiohttpResp( + 200, + { + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/session1", + }, + b"", + ) + resp_412 = _MockAiohttpResp(412, {}, b"Precondition Failed") + query_resp = _MockAiohttpResp( + 200, + {"X-Goog-Upload-Status": "active", "X-Goog-Upload-Size-Received": "100"}, + b"", + ) + final_resp = _MockAiohttpResp(200, {"X-Goog-Upload-Status": "final"}, b"done") + transport.request.side_effect = [start_resp, resp_412, query_resp, final_resp] + + payload = b"A" * 100 + b"B" * 100 + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(chunk_size=100), + transport=transport, + ) + await session.upload(stream=payload) + + transmitted_chunks = [ + request_call.kwargs.get("data") + for request_call in transport.request.call_args_list + if request_call.kwargs.get("headers", {}).get("X-Goog-Upload-Command") + in ("upload", "upload, finalize") + ] + total_bytes_sent = sum( + len(chunk) for chunk in transmitted_chunks if chunk is not None + ) + assert total_bytes_sent == 200, ( + f"Expected 200 bytes uploaded across all chunks, but received {total_bytes_sent} " + f"bytes. Transmitted chunks: {transmitted_chunks}" + ) + + +def test_partial_chunk_recovery_does_not_finalize_prematurely_sync(): + """Verifies that retransmitting a partial chunk does not prematurely finalize the upload. + + Why this test is needed: + When chunk 1 (100 bytes out of a 200-byte stream) partially commits 50 bytes before + failing, `_reposition_stream_offset` slices the active buffer to the remaining 50 bytes. + If EOF detection checks `len(data) < chunk_size` on the sliced buffer instead of + preserving `_buffered_chunk_is_last` computed when reading from the stream, the 50-byte + partial retransmission would be misclassified as the final chunk (`upload, finalize`), + truncating the upload before chunk 2 is ever sent. + """ + transport = mock.Mock() + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/session1", + }, + ) + resp_412 = mock.Mock( + ok=False, + status_code=412, + text="Precondition Failed", + content=b"Precondition Failed", + headers={}, + json=lambda: {}, + ) + # Server confirms 50 of the 100 bytes in chunk 1 were committed. + query_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "50", + }, + ) + final_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b"done", + ) + transport.request.side_effect = [start_resp, resp_412, query_resp, final_resp] + + # Total payload: 200 bytes across two 100-byte chunks. + payload = b"A" * 100 + b"B" * 100 + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(chunk_size=100), + transport=transport, + ) + session.upload(stream=payload) + + # Request sequence: [0] start -> [1] initial chunk 1 -> [2] query -> [3] partial chunk 1 retransmit. + request_calls = transport.request.call_args_list + partial_retransmit_call = request_calls[3] + partial_retransmit_command = partial_retransmit_call.kwargs.get("headers", {}).get( + "X-Goog-Upload-Command" + ) + assert partial_retransmit_command == "upload", ( + f"Partial retransmit request used command '{partial_retransmit_command}' instead of 'upload'. " + "A partial retransmit should not mark finalize when subsequent chunks remain." + ) + + +@pytest.mark.skipif( + not GOOGLE_AUTH_AIO_INSTALLED, + reason="Skipped because google-api-core[async_rest] is not installed", +) +@pytest.mark.asyncio +async def test_partial_chunk_recovery_does_not_finalize_prematurely_async(): + """Async counterpart verifying partial chunk retransmission does not prematurely finalize. + + Why this test is needed: + Ensures ``AsyncResumableUploadSession`` preserves ``_buffered_chunk_is_last`` across + partial recovery so a sliced tail smaller than ``chunk_size`` is sent with + ``X-Goog-Upload-Command: upload`` rather than ``upload, finalize``. + """ + transport = mock.Mock() + start_resp = _MockAiohttpResp( + 200, + { + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/session1", + }, + b"", + ) + resp_412 = _MockAiohttpResp(412, {}, b"Precondition Failed") + query_resp = _MockAiohttpResp( + 200, + {"X-Goog-Upload-Status": "active", "X-Goog-Upload-Size-Received": "50"}, + b"", + ) + final_resp = _MockAiohttpResp(200, {"X-Goog-Upload-Status": "final"}, b"done") + transport.request.side_effect = [start_resp, resp_412, query_resp, final_resp] + + payload = b"A" * 100 + b"B" * 100 + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(chunk_size=100), + transport=transport, + ) + await session.upload(stream=payload) + + # Request sequence: [0] start -> [1] initial chunk 1 -> [2] query -> [3] partial chunk 1 retransmit. + request_calls = transport.request.call_args_list + partial_retransmit_call = request_calls[3] + partial_retransmit_command = partial_retransmit_call.kwargs.get("headers", {}).get( + "X-Goog-Upload-Command" + ) + assert partial_retransmit_command == "upload", ( + f"Async partial retransmit request used command '{partial_retransmit_command}' instead of 'upload'. " + "A partial retransmit should not mark finalize when subsequent chunks remain." + ) + + +# ============================================================================== +# 2. Sync vs Async Consistency & Parity +# ============================================================================== + + +@pytest.mark.skipif( + not GOOGLE_AUTH_AIO_INSTALLED, + reason="Skipped because google-api-core[async_rest] is not installed", +) +@pytest.mark.asyncio +async def test_sync_async_return_type_parity(): + """Verifies consistent return types between sync and async when response_type is set. + + Why this test is needed: + Sync sessions receive a ``requests.Response`` object while async sessions read raw + body bytes inside the ``aiohttp`` context manager. Both must pass the final payload + through ``_format_response_payload`` so callers configuring ``response_type`` + receive identical deserialized types from ``upload()``. + """ + # Sync session + transport_sync = mock.Mock() + start_sync = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/s1", + }, + ) + chunk_sync = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b'{"status": "ok"}', + ) + transport_sync.request.side_effect = [start_sync, chunk_sync] + sync_session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(), + response_type=bytes, + transport=transport_sync, + ) + sync_result = sync_session.upload(stream=b"data") + + # Async session + transport_async = mock.Mock() + start_async = _MockAiohttpResp( + 200, + { + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/s1", + }, + b"", + ) + chunk_async = _MockAiohttpResp( + 200, + {"X-Goog-Upload-Status": "final"}, + b'{"status": "ok"}', + ) + transport_async.request.side_effect = [start_async, chunk_async] + async_session = AsyncResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(), + response_type=bytes, + transport=transport_async, + ) + async_result = await async_session.upload(stream=b"data") + + assert type(sync_result) is type(async_result), ( + f"Expected consistent return types: sync returned {type(sync_result).__name__}, " + f"while async returned {type(async_result).__name__}." + ) + + +@pytest.mark.skipif( + not GOOGLE_AUTH_AIO_INSTALLED, + reason="Skipped because google-api-core[async_rest] is not installed", +) +@pytest.mark.asyncio +async def test_sync_async_custom_retry_parity(): + """Verifies that user-configured retry policies are applied in both sync and async sessions. + + Why this test is needed: + Callers may pass a ``StreamingRetry`` / ``AsyncStreamingRetry`` override to + ``upload(retry=...)`` to customize backoff timing (e.g. ``initial`` delay). + Both sync and async sessions must apply the caller's configured backoff + parameters while wrapping the predicate with protocol recovery rules. + """ + sync_retry = google.api_core.retry.StreamingRetry(initial=0.25) + async_retry = google.api_core.retry.AsyncStreamingRetry(initial=0.25) + + # Sync + transport_sync = mock.Mock() + start_sync = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/s1", + }, + ) + chunk_sync = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b"ok", + ) + transport_sync.request.side_effect = [start_sync, chunk_sync] + sync_session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(), + transport=transport_sync, + ) + resolved_sync = sync_session._get_streaming_retry(retry_override=sync_retry) + sync_session.upload(stream=b"data", retry=sync_retry) + + # Async + transport_async = mock.Mock() + start_async = _MockAiohttpResp( + 200, + { + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/s1", + }, + b"", + ) + chunk_async = _MockAiohttpResp(200, {"X-Goog-Upload-Status": "final"}, b"ok") + transport_async.request.side_effect = [start_async, chunk_async] + async_session = AsyncResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(), + transport=transport_async, + ) + resolved_async = async_session._get_async_streaming_retry( + retry_override=async_retry + ) + await async_session.upload(stream=b"data", retry=async_retry) + + assert resolved_sync._initial == resolved_async._initial == 0.25 + + +@pytest.mark.skipif( + not GOOGLE_AUTH_AIO_INSTALLED, + reason="Skipped because google-api-core[async_rest] is not installed", +) +@pytest.mark.asyncio +async def test_sync_async_error_retry_parity(): + """Verifies that non-retriable exceptions (e.g. ValueError) fail fast without retries in both. + + Why this test is needed: + Programming and configuration errors (such as ``ValueError`` or ``TypeError``) are + not transient transport failures. Both sync and async retry predicates must reject + non-retriable exceptions immediately on the first attempt without entering backoff loops. + """ + # Sync: transport raises ValueError + transport_sync = mock.Mock() + transport_sync.request.side_effect = ValueError("fatal configuration issue") + sync_session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(), + transport=transport_sync, + ) + with pytest.raises(ValueError): + sync_session.upload(stream=b"data") + sync_attempts = transport_sync.request.call_count + + # Async: transport raises ValueError + transport_async = mock.Mock() + transport_async.request.side_effect = ValueError("fatal configuration issue") + async_session = AsyncResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(), + transport=transport_async, + ) + with pytest.raises(ValueError): + await async_session.upload(stream=b"data") + async_attempts = transport_async.request.call_count + + assert sync_attempts == async_attempts == 1, ( + f"Expected non-retriable error to fail after 1 attempt in both implementations: " + f"sync attempts={sync_attempts}, async attempts={async_attempts}." + ) + + +def test_iterable_stream_incremental_consumption_sync(): + """Verifies that iterable streams are consumed incrementally rather than loaded upfront. + + Why this test is needed: + Resumable uploads often stream multi-gigabyte payloads from generators or chunk + iterators. If ``_prepare_stream`` eagerly materialized the iterable into memory + (e.g. via ``b"".join(stream)``) to inspect its size, large uploads would exhaust + process memory before transmitting the first chunk. + """ + chunks_generated = 0 + + def data_generator(): + nonlocal chunks_generated + for _ in range(10): + chunks_generated += 1 + yield b"X" * 1024 + + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(chunk_size=1024), + ) + # Stream preparation must wrap the generator lazily without advancing it. + stream_obj, size = session._prepare_stream(data_generator(), size=10 * 1024) + + assert chunks_generated == 0, ( + f"Expected iterable stream to be consumed lazily during upload, but all " + f"{chunks_generated} chunks were buffered into memory during preparation." + ) + + +def test_timezone_naive_deadline_handling(): + """Verifies that naive future datetimes are compared accurately against current time. + + Why this test is needed: + Callers frequently construct upload deadlines using ``datetime.datetime.now()`` + (timezone-naive) rather than ``datetime.datetime.now(datetime.timezone.utc)``. + Comparing a naive ``deadline`` against a timezone-aware clock raises ``TypeError``, + and misinterpreting local time as UTC skews the remaining deadline budget. + """ + future_naive = datetime.datetime.now() + datetime.timedelta(hours=1) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(deadline=future_naive), + ) + remaining = session._get_deadline_remaining() + assert remaining is not None and remaining > 0, ( + f"Expected positive remaining deadline duration for future naive timestamp, got {remaining}s." + ) + + +# ============================================================================== +# 3. Retry Coordination & Multi-Layer Interactions +# ============================================================================== + + +def test_single_layer_backoff_coordination(): + """Verifies that retry backoff delays are coordinated through a single layer during chunk recovery. + + Why this test is needed: + When a chunk upload fails with a transient 503 error, recovery involves querying + the server offset before retransmitting. If inner helper methods wrapped their own + independent retry loops around chunk transmission and recovery, a single transient + failure would trigger compounded backoff sleeps across multiple layers. Only the + outer streaming retry loop should schedule backoff delays. + """ + transport = mock.Mock() + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/s1", + }, + ) + resp_503 = mock.Mock( + ok=False, + status_code=503, + text="Service Unavailable", + content=b"503", + headers={}, + json=lambda: {}, + ) + query_200 = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "0", + }, + ) + final_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b"done", + ) + + transport.request.side_effect = [ + start_resp, + resp_503, + query_200, + final_resp, + ] + + sleep_durations = [] + with mock.patch( + "time.sleep", side_effect=lambda duration: sleep_durations.append(duration) + ): + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(chunk_size=100), + transport=transport, + ) + session.upload(stream=b"A" * 100) + + # For a single transient failure, exactly 1 coordinated retry backoff sleep should occur. + assert len(sleep_durations) == 1, ( + f"Expected exactly 1 coordinated retry backoff delay for a single failed chunk attempt, " + f"but observed {len(sleep_durations)} sleeps ({sleep_durations}) compounding across nested layers." + ) + + +def test_connection_recovery_with_custom_retry_predicate(): + """Verifies that network connection drops initiate an offset query when a custom retry is set. + + Why this test is needed: + Callers often supply a narrow custom retry predicate (such as retrying only + ``ServiceUnavailable``) to control transient HTTP status retries. When a TCP + connection resets mid-chunk, the client cannot tell how many bytes the server + persisted and must issue an ``X-Goog-Upload-Command: query`` offset check. + Transport connection drops must remain recoverable even when a custom predicate is set. + """ + transport = mock.Mock() + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/s1", + }, + ) + query_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "0", + }, + ) + final_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b"done", + ) + + transport.request.side_effect = [ + start_resp, + requests.exceptions.ConnectionError("TCP connection reset by peer"), + query_resp, + final_resp, + ] + + user_retry = google.api_core.retry.StreamingRetry( + predicate=google.api_core.retry.if_exception_type( + google.api_core.exceptions.ServiceUnavailable + ) + ) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(chunk_size=100), + transport=transport, + ) + + session.upload(stream=b"A" * 100, retry=user_retry) + + request_calls = transport.request.call_args_list + query_calls = [ + request_call + for request_call in request_calls + if request_call.kwargs.get("headers", {}).get("X-Goog-Upload-Command") + == "query" + ] + assert len(query_calls) >= 1, ( + "Expected transport connection drop to initiate an offset query to reconcile server state." + ) + + +def test_precondition_failed_response_triggers_offset_query(): + """Verifies that a 412 Precondition Failed status initiates an offset query rather than immediate resend. + + Why this test is needed: + Under the Resumable Upload protocol, a ``412 Precondition Failed`` response during + chunk upload indicates a state/offset mismatch between client and server. Blindly + retransmitting the chunk at the old offset without first issuing an + ``X-Goog-Upload-Command: query`` request would repeatedly fail with 412. The session + must set ``_needs_recovery = True`` so the immediate next request queries server state. + """ + transport = mock.Mock() + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/s1", + }, + ) + resp_412 = mock.Mock( + ok=False, + status_code=412, + text="Precondition Failed", + content=b"Precondition Failed", + headers={}, + json=lambda: {}, + ) + query_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "0", + }, + ) + final_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b"done", + ) + + transport.request.side_effect = [ + start_resp, + resp_412, + query_resp, + final_resp, + ] + + user_retry = google.api_core.retry.StreamingRetry( + maximum=0.01, + deadline=0.05, + ) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig(chunk_size=100), + transport=transport, + ) + + session.upload(stream=b"A" * 100, retry=user_retry) + + # Request sequence: [0] start -> [1] failed chunk upload (412) -> [2] recovery query. + request_calls = transport.request.call_args_list + recovery_request_call = request_calls[2] + recovery_request_command = recovery_request_call.kwargs.get("headers", {}).get( + "X-Goog-Upload-Command" + ) + assert recovery_request_command == "query", ( + f"Expected command 'query' to reconcile server offset after 412 Precondition Failed, " + f"but received '{recovery_request_command}'." + ) + + +def test_socket_timeout_initiates_recovery(): + """Verifies that a transient socket read timeout triggers recovery when stall termination is disabled. + + Why this test is needed: + A socket read timeout (``requests.exceptions.Timeout``) during chunk upload means + the client cannot determine how many bytes the server persisted before the connection + stalled. When stall termination is disabled (``stall_minimum_rate=0``), the session + must treat the timeout as a recoverable transport interruption, query the server + offset via ``X-Goog-Upload-Command: query``, and resume uploading from that offset. + """ + transport = mock.Mock() + start_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-URL": "https://upload.example.com/s1", + }, + ) + query_resp = mock.Mock( + ok=True, + status_code=200, + headers={ + "X-Goog-Upload-Status": "active", + "X-Goog-Upload-Size-Received": "0", + }, + ) + final_resp = mock.Mock( + ok=True, + status_code=200, + headers={"X-Goog-Upload-Status": "final"}, + content=b"done", + ) + transport.request.side_effect = [ + start_resp, + requests.exceptions.Timeout("Read timed out on socket"), + query_resp, + final_resp, + ] + + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=ResumableUploadConfig( + chunk_size=100, stall_minimum_rate=0, stall_timeout=0.0 + ), + transport=transport, + ) + + session.upload(stream=b"A" * 100) + + request_calls = transport.request.call_args_list + assert len(request_calls) > 2, ( + f"Expected socket read timeout to attempt recovery rather than terminating after {len(request_calls)} call(s)." + ) + + +@pytest.mark.skipif( + not GOOGLE_AUTH_AIO_INSTALLED, + reason="Skipped because google-api-core[async_rest] is not installed", +) +@pytest.mark.asyncio +async def test_async_upload_cancellation_does_not_deadlock(): + """Verifies that cancelling an in-flight async upload terminates cleanly without queue deadlock. + + Why this test is needed: + ``AsyncUploadOperation`` runs the upload in a background ``asyncio.Task`` and feeds + progress snapshots to callers iterating over ``op.progress()`` via an internal + ``asyncio.Queue``. In Python 3.8+, ``asyncio.CancelledError`` inherits from + ``BaseException`` rather than ``Exception``. If the background task catches only + ``Exception`` when pushing error sentinels into ``_progress_queue``, cancelling + the task leaves ``op.progress()`` consumers awaiting ``_progress_queue.get()`` + forever. This test ensures task cancellation propagates cleanly to consumers. + """ + + class HangingTransport: + def __init__(self): + self.started = asyncio.Event() + + def request(self, method, url, **kwargs): + return self._HangingRequestContext(self) + + class _HangingRequestContext: + def __init__(self, parent): + self.parent = parent + + async def __aenter__(self): + self.parent.started.set() + await asyncio.sleep(60) + + async def __aexit__(self, *args): + pass + + transport = HangingTransport() + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/init", + transport=transport, + ) + + upload_operation = session.upload(stream=b"payload") + progress_iterator = upload_operation.progress() + consume_task = asyncio.create_task(progress_iterator.__anext__()) + await transport.started.wait() + consume_task.cancel() + with pytest.raises((asyncio.CancelledError, StopAsyncIteration)): + await asyncio.wait_for(consume_task, timeout=5)