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..041ef2a6cf72 --- /dev/null +++ b/packages/google-api-core/google/api_core/resumable_transfer/__init__.py @@ -0,0 +1,55 @@ +# 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, + Command, + ProgressState, + Status, + UploadProgress, +) +from google.api_core.resumable_transfer.upload import ( + ResumableUploadConfig, + ResumableUploadSession, +) +from google.api_core.resumable_transfer.upload_async import ( + AsyncResumableUploadSession, + AsyncUploadOperation, +) + +__all__ = [ + "Command", + "DEFAULT_CHUNK_SIZE", + "MissingStatusHeaderError", + "ProgressState", + "ResumableTransferError", + "Status", + "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..0859cde83de3 --- /dev/null +++ b/packages/google-api-core/google/api_core/resumable_transfer/common.py @@ -0,0 +1,89 @@ +# 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 and headers for Resumable Upload protocol.""" + +import dataclasses +import enum +from typing import Optional + +# Default chunk size: 10 MiB +DEFAULT_CHUNK_SIZE = 10 * 1024 * 1024 + +# 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 +RETRYABLE_STATUS_CODES = (408, 429, 500, 502, 503, 504) + +# HTTP status codes indicating state consistency errors requiring recovery +RECOVERABLE_STATUS_CODES = (400, 409, 412, 416) 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..39089d40eb17 --- /dev/null +++ b/packages/google-api-core/google/api_core/resumable_transfer/upload.py @@ -0,0 +1,912 @@ +# 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 contextlib +import dataclasses +import datetime +import io +import logging +import time +from typing import ( + Any, + BinaryIO, + Callable, + Generator, + Iterable, + List, + Mapping, + Optional, + Sequence, + Tuple, + Union, + cast, +) + +import google.protobuf.message +import proto +import requests +from google.protobuf import json_format + +import google.api_core.retry +from google.api_core import exceptions +from google.api_core.resumable_transfer import common, upload_state + +_LOGGER = logging.getLogger(__name__) +_DEFAULT_START_TIMEOUT = 60.0 # seconds for initial start request +_monotonic_clock = time.monotonic + + +class _RecoveryRetransmit(Exception): + """Internal exception indicating state synchronization succeeded and chunk should retransmit.""" + + pass + + +@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. + start_timeout: Local per-request timeout in seconds for start request. + start_retry: Custom retry policy for the start request. + 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: Overall global deadline for the upload process. + timeout: Fallback per-request timeout. + retry: Fallback retry policy. + on_progress: Callback function receiving UploadProgress notifications. + response_type: Optional message class (proto.Message or google.protobuf.message.Message), + callable deserializer, or None to return raw response. + content_type: MIME type of the stream payload. + """ + + chunk_size: int = common.DEFAULT_CHUNK_SIZE + start_timeout: Optional[float] = None + start_retry: Optional[google.api_core.retry.Retry] = None + 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 + timeout: Optional[float] = None + retry: Optional[google.api_core.retry.Retry] = None + on_progress: Optional[Callable[[common.UploadProgress], None]] = None + response_type: Optional[Any] = None + content_type: Optional[str] = None + + def __post_init__(self) -> None: + """Normalizes fallback timeouts and retry policies.""" + if self.start_timeout is not None and self.timeout is None: + self.timeout = self.start_timeout + elif self.timeout is not None and self.start_timeout is None: + self.start_timeout = self.timeout + + if self.start_retry is not None and self.retry is None: + self.retry = self.start_retry + elif self.retry is not None and self.start_retry is None: + self.start_retry = self.retry + + @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) + + +class ResumableUploadSession: + """Manages the full lifecycle of a resumable upload session.""" + + def __init__( + self, + upload_url: Optional[str] = None, + config: Optional[ResumableUploadConfig] = None, + resumable_url: Optional[str] = None, + transport: Optional[requests.Session] = None, + ) -> None: + """Initializes a ResumableUploadSession. + + Args: + upload_url: The initial URL for the start request. + config: Optional upload configuration parameters. + resumable_url: Pre-existing upload session URL if resuming. + transport: Optional requests session. + """ + self._config = config or ResumableUploadConfig() + self._transport = transport + self._response: Optional[Any] = None + self._state = upload_state.ProtocolState( + upload_url=upload_url, + chunk_size=self._config.chunk_size, + resumable_url=resumable_url, + ) + + # In-memory zero-copy buffer (never discard chunk until confirmed) + self._buffered_chunk: Optional[memoryview] = None + self._buffered_chunk_offset: int = 0 + 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._captured_progress: Optional[List[common.UploadProgress]] = None + + @property + def upload_url(self) -> Optional[str]: + """Optional[str]: The unique upload URL for this session.""" + return self._state.resumable_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 _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. + """ + if hasattr(exc, "__dict__"): + setattr(exc, "upload_url", self.upload_url) + setattr(exc, "chunk_size", self.chunk_size) + + def _notify_progress(self, state: common.ProgressState) -> None: + """Notifies progress with current upload status. + + Args: + state: ProgressState transition milestone. + """ + if self.upload_url: + progress = common.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 self._captured_progress is not None: + self._captured_progress.append(progress) + if self._config.on_progress: + self._config.on_progress(progress) + + @contextlib.contextmanager + def _capture_progress( + self, + ) -> Generator[List[common.UploadProgress], None, None]: + """Intercepts progress events to buffer snapshots for generator consumers. + + Yields: + List buffering UploadProgress snapshots during generator execution. + """ + captured: List[common.UploadProgress] = [] + self._captured_progress = captured + try: + yield captured + finally: + self._captured_progress = None + + 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 + if dl.tzinfo is None: + dl = dl.replace(tzinfo=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) -> float: + """Computes timeout in seconds for start and control requests. + + Returns: + Applicable timeout in seconds. + """ + remaining = self._get_deadline_remaining() + timeout = ( + self._config.start_timeout or self._config.timeout or _DEFAULT_START_TIMEOUT + ) + if remaining is not None: + return min(timeout, remaining) + return timeout + + def _get_retry_predicate(self) -> Callable[[Any], bool]: + """Returns a predicate function for determining if an exception is retryable. + + Returns: + A callable accepting an exception and returning a boolean. + """ + + def should_retry(exc: Any) -> bool: + if isinstance( + exc, + ( + exceptions.DeadlineExceeded, + exceptions.TransferStalledError, + exceptions.UploadCancelledError, + ), + ): + return False + if isinstance(exc, exceptions.MissingStatusHeaderError): + return True + if isinstance(exc, requests.exceptions.RequestException): + if isinstance( + exc, + ( + requests.exceptions.ConnectionError, + requests.exceptions.ChunkedEncodingError, + requests.exceptions.Timeout, + ), + ): + return True + if isinstance(exc, exceptions.GoogleAPICallError): + return exc.code in common.RETRYABLE_STATUS_CODES + return False + + return should_retry + + def _get_retry(self, is_start: bool = False) -> google.api_core.retry.Retry: + """Resolves retry policy for requests. + + Args: + is_start: Whether this retry policy is for the start request. + + Returns: + Configured or default Retry instance. + """ + if is_start and self._config.start_retry: + return self._config.start_retry + if self._config.retry: + return self._config.retry + return google.api_core.retry.Retry(predicate=self._get_retry_predicate()) + + def _compute_chunk_timeout(self, data_len: int) -> float: + """Computes the dynamic per-attempt chunk timeout based on stall control and deadlines. + + Args: + data_len: Length of the current chunk in bytes. + + 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 + next_chunk_timeout = max( + 1.0, + expected_sec - self._aggregate_lag + self._config.stall_timeout, + ) + per_attempt_timeout = max(5.0, min(next_chunk_timeout, 2.0 * expected_sec)) + + if self._config.timeout: + per_attempt_timeout = min(self._config.timeout, per_attempt_timeout) + + 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_start: float, 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_start: Monotonic timestamp before chunk transmission began. + 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: + if self._stall_timeout_started is None: + self._stall_timeout_started = t_start + if ( + _monotonic_clock() - 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: BinaryIO, 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 + return received + + self._buffered_chunk = None + if hasattr(stream, "seekable") and not stream.seekable(): + 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: + stream.seek(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, + ) -> 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. + + Returns: + The upload session URL. + """ + method, url, headers, payload = self._state.build_start_request( + body=request_body, + headers=self._config.start_headers, + content_type=self._config.content_type, + size=size, + ) + + def do_initiate() -> str: + timeout = self._get_start_timeout() + response = transport.request( + method, url, data=payload, headers=headers, timeout=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 + + session_url = self._get_retry(is_start=True)(do_initiate)() + self._notify_progress(common.ProgressState.STARTED) + return session_url + + def _transmit_chunk( + self, transport: requests.Session, stream: BinaryIO, size: Optional[int] + ) -> requests.Response: + """Transmits the next data chunk with stall control and error recovery. + + Args: + transport: The requests session. + stream: The input data stream. + size: Total size of the stream in bytes, if known. + + Returns: + The HTTP response for the transmitted chunk. + """ + + def do_transmit() -> requests.Response: + chunk_size = self._state.chunk_size + + # Retain active chunk in zero-copy buffer if not present + if self._buffered_chunk is None: + raw_bytes = stream.read(chunk_size) + if not raw_bytes: + raw_bytes = b"" + self._buffered_chunk = memoryview(raw_bytes) + self._buffered_chunk_offset = self._state.bytes_uploaded + + data = self._buffered_chunk + data_len = len(data) + + is_last = data_len < chunk_size + if size is not None and self._state.bytes_uploaded + data_len >= size: + is_last = True + + method, url, headers, payload = self._state.build_chunk_request( + data=data, + is_last_chunk=is_last, + content_type=self._config.content_type, + ) + + def do_http() -> requests.Response: + per_attempt_timeout = self._compute_chunk_timeout(data_len) + try: + resp = transport.request( + method, + url, + data=payload, + headers=headers, + timeout=per_attempt_timeout, + ) + except requests.exceptions.Timeout as exc: + if self._config.stall_minimum_rate and self._config.stall_timeout: + remaining = self._get_deadline_remaining() + if remaining is not None and remaining <= 0: + raise exceptions.DeadlineExceeded( + f"Resumable upload deadline {self._config.deadline} exceeded." + ) from exc + raise exceptions.TransferStalledError( + f"Upload stalled: chunk transfer timed out ({exc}).", + upload_url=self.upload_url, + chunk_size=self.chunk_size, + ) from exc + raise + if not resp.ok: + raise exceptions.from_http_response(resp) + return resp + + try: + t_start = _monotonic_clock() + resp = self._get_retry()(do_http)() + t_elapsed = _monotonic_clock() - t_start + + self._update_stall_control(data_len, t_start, 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 + ) + return resp + except Exception as exc: + self._enrich_exception(exc) + if isinstance(exc, exceptions.DeadlineExceeded): + raise + if isinstance(exc, requests.exceptions.Timeout): + remaining = self._get_deadline_remaining() + if remaining is not None and remaining <= 0: + raise exceptions.DeadlineExceeded( + f"Resumable upload deadline {self._config.deadline} exceeded." + ) from exc + raise exceptions.TransferStalledError( + f"Upload stalled: chunk transfer timed out ({exc}).", + upload_url=self.upload_url, + chunk_size=self.chunk_size, + ) from exc + + is_recoverable = ( + isinstance(exc, exceptions.GoogleAPICallError) + and exc.code in common.RECOVERABLE_STATUS_CODES + ) or isinstance(exc, exceptions.MissingStatusHeaderError) + + if is_recoverable: + _LOGGER.info( + "Recoverable error %s during chunk upload. Querying server offset.", + exc, + ) + self._notify_progress(common.ProgressState.RECOVERING) + self._recover(transport, stream) + raise _RecoveryRetransmit() + raise + + recovery_loop = google.api_core.retry.Retry( + predicate=lambda e: isinstance(e, _RecoveryRetransmit) + ) + return recovery_loop(do_transmit)() + + def _recover(self, transport: requests.Session, stream: BinaryIO) -> int: + """Queries server for committed byte offset and adjusts buffer / stream. + + Args: + transport: The requests session. + stream: The input data stream. + + Returns: + The confirmed server byte offset. + + 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() + + def do_query() -> requests.Response: + 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) + return resp + + resp = self._get_retry()(do_query)() + received = self._state.process_query_response(resp.status_code, resp.headers) + self._notify_progress(common.ProgressState.OFFSET_RECEIVED) + return self._reposition_stream_offset(stream, received) + + 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: BinaryIO, + computed_size: Optional[int], + captured: Optional[List[common.UploadProgress]] = None, + ) -> Generator[common.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. + captured: Optional buffer accumulating progress snapshots. + + Yields: + UploadProgress snapshots for each transmission milestone. + + Raises: + ValueError: If upload concludes without a server response. + """ + if captured: + while captured: + yield captured.pop(0) + + final_resp = None + while not self._state.finished and not self._state.invalid: + final_resp = self._transmit_chunk(transport, stream_obj, computed_size) + if captured: + while captured: + yield captured.pop(0) + + if final_resp is None: + raise ValueError("Upload completed without receiving a final response.") + + self._response = self._format_response(final_resp) + + def upload( + self, + stream: Union[BinaryIO, bytes, Iterable[bytes]], + request_body: Union[str, bytes] = "", + size: Optional[int] = None, + transport: Optional[requests.Session] = 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. + + 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 + ): + 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, + ) -> Generator[common.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. + + 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) + with self._capture_progress() as captured: + try: + stream_obj, computed_size = self._prepare_stream(stream, size) + self.initiate( + transport=sess, request_body=request_body, size=computed_size + ) + yield from self._transmit_all_chunks( + sess, stream_obj, computed_size, captured + ) + except Exception as exc: + self._enrich_exception(exc) + raise + + def resume( + self, + upload_url: Optional[str] = None, + stream: Optional[Union[BinaryIO, bytes, Iterable[bytes]]] = None, + size: Optional[int] = None, + chunk_size: Optional[int] = None, + transport: Optional[requests.Session] = None, + ) -> Any: + """Resumes an existing upload from a saved upload URL. + + Args: + upload_url: The pre-existing upload session URL. + stream: The data payload to resume uploading from. + size: Total size of the payload in bytes, if known. + chunk_size: Optional chunk size override in bytes. + transport: Optional requests session. + + 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, + ): + pass + return self._response + + def iter_resume( + self, + upload_url: Optional[str] = None, + stream: Optional[Union[BinaryIO, bytes, Iterable[bytes]]] = None, + size: Optional[int] = None, + chunk_size: Optional[int] = None, + transport: Optional[requests.Session] = None, + ) -> Generator[common.UploadProgress, None, None]: + """Streams resumption of an upload, yielding UploadProgress snapshots. + + Args: + upload_url: The pre-existing upload session URL. + stream: The data payload to resume uploading from. + size: Total size of the payload in bytes, if known. + chunk_size: Optional chunk size override in bytes. + transport: Optional requests session. + + 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.") + + if chunk_size is not None: + self._state._chunk_size = chunk_size + + self._state._resumable_url = actual_url + with self._capture_progress() as captured: + try: + stream_obj, computed_size = self._prepare_stream(stream, size) + self._recover(sess, stream_obj) + yield from self._transmit_all_chunks( + sess, stream_obj, computed_size, captured + ) + except Exception as exc: + self._enrich_exception(exc) + raise + + def _prepare_stream( + self, stream: Union[BinaryIO, bytes, Iterable[bytes]], size: Optional[int] + ) -> Tuple[BinaryIO, Optional[int]]: + """Normalizes stream input into a BinaryIO 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 BinaryIO stream, 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: BinaryIO = io.BytesIO(stream) + if computed_size is None: + computed_size = len(stream) + elif not hasattr(stream, "read") and isinstance(stream, Iterable): + stream_obj = io.BytesIO(b"".join(stream)) + if computed_size is None: + computed_size = stream_obj.getbuffer().nbytes + elif hasattr(stream, "read"): + stream_obj = cast(BinaryIO, stream) + if computed_size is None: + if hasattr(stream_obj, "getbuffer"): + computed_size = stream_obj.getbuffer().nbytes + elif ( + hasattr(stream_obj, "seekable") + and stream_obj.seekable() + and hasattr(stream_obj, "tell") + ): + cur = stream_obj.tell() + stream_obj.seek(0, io.SEEK_END) + computed_size = stream_obj.tell() - cur + stream_obj.seek(cur) + else: + raise TypeError(f"Unsupported stream type: {type(stream)}") + + if hasattr(stream_obj, "tell"): + try: + self._start_stream_offset = stream_obj.tell() + except (OSError, AttributeError): + self._start_stream_offset = 0 + + return stream_obj, computed_size + + def _format_response(self, response: requests.Response) -> Any: + """Formats response into protobuf message type if provided. + + Args: + response: HTTP response object from final chunk. + + Returns: + Deserialized protobuf message or the raw response object. + """ + return _format_response_payload(response, self._config.response_type) + + +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 the raw response object / bytes. + """ + if response_type is None: + return response + + content: bytes + if isinstance(response, bytes): + content = response + elif hasattr(response, "content"): + content = response.content + else: + content = bytes(response) + + if isinstance(response_type, type) and issubclass(response_type, proto.Message): + return cast(Any, response_type).from_json(content, ignore_unknown_fields=True) + 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 hasattr(response_type, "from_json") and callable(response_type.from_json): + return response_type.from_json(content) + if callable(response_type): + return response_type(content) + + return response 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..cba5ad3e5ca9 --- /dev/null +++ b/packages/google-api-core/google/api_core/resumable_transfer/upload_async.py @@ -0,0 +1,877 @@ +# 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, + AsyncIterable, + AsyncIterator, + Awaitable, + BinaryIO, + Callable, + Generator, + Generic, + Iterable, + Mapping, + Optional, + Tuple, + TypeVar, + Union, +) + +try: + import aiohttp +except ImportError: # pragma: NO COVER + aiohttp = None # type: ignore + +from google.api_core import exceptions +from google.api_core.resumable_transfer import common, upload_state +from google.api_core.resumable_transfer.upload import ( + ResumableUploadConfig, + _format_response_payload, +) + +_LOGGER = logging.getLogger(__name__) +_DEFAULT_START_TIMEOUT = 60.0 # seconds for initial start request +_DONE_SENTINEL = object() +_monotonic_clock = time.monotonic + +ResponseProto = TypeVar("ResponseProto") + + +class _AsyncRecoveryRetransmit(Exception): + """Internal exception indicating state synchronization succeeded and chunk should retransmit.""" + + pass + + +class AsyncUploadOperation(Generic[ResponseProto], Awaitable[ResponseProto]): + """Handle representing an active asynchronous upload operation. + + Implements Awaitable[ResponseProto] so awaiting the operation directly + returns the deserialized response upon transfer completion. + """ + + def __init__( + self, + task: asyncio.Task, + session: "AsyncResumableUploadSession", + progress_queue: asyncio.Queue, + ) -> None: + """Initializes the active upload operation handle. + + Args: + task: Background asyncio task driving the upload. + session: Underlying asynchronous resumable upload session. + progress_queue: Queue used to deliver upload progress updates. + """ + self._task = task + self._session = session + self._progress_queue = progress_queue + + def __await__(self) -> Generator[Any, None, ResponseProto]: + """Awaits completion of the upload task and returns the server response.""" + return self._task.__await__() + + async def progress(self) -> AsyncIterator[common.UploadProgress]: + """Returns an asynchronous stream yielding progress snapshots without blocking uploads. + + Yields: + UploadProgress snapshots for each progress transition. + + Raises: + Exception: Re-raises any exception encountered during the background transfer. + """ + while True: + item = await self._progress_queue.get() + if item is _DONE_SENTINEL: + break + if isinstance(item, Exception): + raise item + yield item + + @property + def response(self) -> Optional[ResponseProto]: + """The deserialized protobuf response message, 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, + resumable_url: Optional[str] = None, + transport: Optional[Any] = None, + ) -> None: + """Initializes an AsyncResumableUploadSession. + + Args: + upload_url: The initial URL for the start request. + config: Optional upload configuration parameters. + resumable_url: Pre-existing upload session URL if resuming. + transport: Optional aiohttp.ClientSession. + """ + self._config = config or ResumableUploadConfig() + self._transport = transport + self._response: Optional[Any] = None + self._state = upload_state.ProtocolState( + upload_url=upload_url, + chunk_size=self._config.chunk_size, + resumable_url=resumable_url, + ) + + # In-memory zero-copy buffer + self._buffered_chunk: Optional[memoryview] = None + self._buffered_chunk_offset: int = 0 + 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 + + @property + def upload_url(self) -> Optional[str]: + """Optional[str]: The unique upload URL for this session.""" + return self._state.resumable_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 _ensure_aiohttp(self) -> None: + """Validates that aiohttp is installed and accessible. + + Raises: + ImportError: If aiohttp is not installed. + """ + if aiohttp is None: + 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. + """ + if hasattr(exc, "__dict__"): + setattr(exc, "upload_url", self.upload_url) + setattr(exc, "chunk_size", self.chunk_size) + + def _notify_progress( + self, state: common.ProgressState, queue: Optional[asyncio.Queue] = None + ) -> None: + """Notifies registered progress callback and queue with current upload status. + + Args: + state: ProgressState transition milestone. + queue: Optional queue to receive progress event. + """ + if self.upload_url: + progress = common.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 self._config.on_progress: + try: + self._config.on_progress(progress) + except Exception: # pragma: NO COVER + pass + if queue is not None: + queue.put_nowait(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 + if dl.tzinfo is None: + dl = dl.replace(tzinfo=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) -> float: + """Computes timeout in seconds for start and control requests. + + Returns: + Applicable timeout in seconds. + """ + remaining = self._get_deadline_remaining() + timeout = ( + self._config.start_timeout or self._config.timeout or _DEFAULT_START_TIMEOUT + ) + if remaining is not None: + return min(timeout, remaining) + return timeout + + async def _async_retry( + self, coro_fn: Callable[[], Awaitable[Any]], max_attempts: int = 4 + ) -> Any: + """Executes an asynchronous callable with exponential backoff retry logic. + + Args: + coro_fn: Asynchronous nullary function to invoke and retry. + max_attempts: Maximum retry attempts before propagating failure. + + Returns: + The successful return value of coro_fn. + """ + delay = 1.0 + multiplier = 2.0 + max_delay = 60.0 + for attempt in range(max_attempts): + try: + return await coro_fn() + except ( + exceptions.DeadlineExceeded, + exceptions.TransferStalledError, + exceptions.UploadCancelledError, + ): + raise + except exceptions.MissingStatusHeaderError: + if attempt == max_attempts - 1: + raise + except exceptions.GoogleAPICallError as exc: + if exc.code not in common.RETRYABLE_STATUS_CODES: + raise + if attempt == max_attempts - 1: + raise + except Exception: + if attempt == max_attempts - 1: + raise + + await asyncio.sleep(delay) + delay = min(delay * multiplier, max_delay) + + async def initiate( + self, + transport: Any, + request_body: Union[str, bytes] = "", + size: Optional[int] = None, + progress_queue: Optional[asyncio.Queue] = None, + ) -> str: + """Initiates the upload session by sending the start command asynchronously. + + Args: + transport: The aiohttp client session. + request_body: Initial metadata payload sent with start command. + size: Total stream size in bytes, if known. + progress_queue: Optional queue to receive progress event. + + Returns: + The upload session URL. + + Raises: + GoogleAPICallError: If the server rejects the start request. + MissingStatusHeaderError: If the server response lacks status header. + """ + self._ensure_aiohttp() + method, url, headers, payload = self._state.build_start_request( + body=request_body, + headers=self._config.start_headers, + content_type=self._config.content_type, + size=size, + ) + + async def do_initiate(): + 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 + + session_url = await self._async_retry(do_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[asyncio.Queue] = None, + stream_obj: Any = 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. + stream_obj: Underlying stream object for recovery seeking. + + 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 unrecoverable error. + """ + + async def do_transmit(): + chunk_size = self._state.chunk_size + + # Retain active chunk in zero-copy buffer if not present + 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 + + data = self._buffered_chunk + data_len = len(data) + + is_last = data_len < chunk_size + if size is not None and self._state.bytes_uploaded + data_len >= size: + is_last = True + + method, url, headers, payload = self._state.build_chunk_request( + data=data, + is_last_chunk=is_last, + content_type=self._config.content_type, + ) + + async def do_http(): + rate = self._config.stall_minimum_rate + expected_sec = data_len / rate if rate > 0 else 60.0 + next_chunk_timeout = max( + 1.0, + expected_sec - self._aggregate_lag + self._config.stall_timeout, + ) + per_attempt_timeout = max( + 5.0, min(next_chunk_timeout, 2.0 * expected_sec) + ) + + if self._config.timeout: + per_attempt_timeout = min(self._config.timeout, per_attempt_timeout) + + remaining = self._get_deadline_remaining() + if remaining is not None: + per_attempt_timeout = min(per_attempt_timeout, remaining) + + client_timeout = aiohttp.ClientTimeout(total=per_attempt_timeout) + 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") + ) + return resp.status, resp_headers, resp_body + except asyncio.TimeoutError as exc: + if self._config.stall_minimum_rate and self._config.stall_timeout: + remaining = self._get_deadline_remaining() + if remaining is not None and remaining <= 0: + raise exceptions.DeadlineExceeded( + f"Resumable upload deadline {self._config.deadline} exceeded." + ) from exc + raise exceptions.TransferStalledError( + f"Upload stalled: chunk transfer timed out ({exc}).", + upload_url=self.upload_url, + chunk_size=self.chunk_size, + ) from exc + raise + + try: + t_start = _monotonic_clock() + status_code, resp_headers, resp_body = await self._async_retry(do_http) + t_elapsed = _monotonic_clock() - t_start + + # Evaluate stall control lag & timer + if self._config.stall_minimum_rate and self._config.stall_timeout: + 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: + if self._stall_timeout_started is None: + self._stall_timeout_started = t_start + if ( + _monotonic_clock() - 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 + + 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 + except Exception as exc: + self._enrich_exception(exc) + if isinstance(exc, exceptions.DeadlineExceeded): + raise + if isinstance( + exc, + ( + asyncio.TimeoutError, + aiohttp.ServerTimeoutError if aiohttp else (), + ), + ): + remaining = self._get_deadline_remaining() + if remaining is not None and remaining <= 0: + raise exceptions.DeadlineExceeded( + f"Resumable upload deadline {self._config.deadline} exceeded." + ) from exc + raise exceptions.TransferStalledError( + f"Upload stalled: chunk transfer timed out ({exc}).", + upload_url=self.upload_url, + chunk_size=self.chunk_size, + ) from exc + + is_recoverable = ( + isinstance(exc, exceptions.GoogleAPICallError) + and exc.code in common.RECOVERABLE_STATUS_CODES + ) or isinstance(exc, exceptions.MissingStatusHeaderError) + + if is_recoverable: + _LOGGER.info( + "Recoverable error %s during async chunk upload. Querying server offset.", + exc, + ) + self._notify_progress( + common.ProgressState.RECOVERING, progress_queue + ) + await self._recover(transport, stream_obj) + raise _AsyncRecoveryRetransmit() + raise + + while True: + try: + return await do_transmit() + except _AsyncRecoveryRetransmit: + continue + + async def _recover(self, transport: Any, stream_obj: Any = None) -> int: + """Queries server for committed byte offset and adjusts buffer. + + Args: + transport: The aiohttp client session. + stream_obj: Underlying stream object to rewind if seekable. + + Returns: + The confirmed server byte offset. + + 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() + + async def do_query(): + 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") + ) + return resp.status, resp_headers + + status_code, resp_headers = await self._async_retry(do_query) + received = self._state.process_query_response(status_code, resp_headers) + + 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 + return received + + self._buffered_chunk = None + if stream_obj is not None and hasattr(stream_obj, "seek"): + if hasattr(stream_obj, "seekable") and not stream_obj.seekable(): + 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: + stream_obj.seek(self._start_stream_offset + received) + return 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 + + 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) + + def upload( + self, + stream: Union[AsyncIterable[bytes], BinaryIO, bytes, Iterable[bytes]], + request_body: Union[str, bytes] = "", + size: Optional[int] = None, + transport: Optional[Any] = 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. + + Returns: + An AsyncUploadOperation handle representing the active 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.") + + progress_queue: asyncio.Queue = asyncio.Queue() + reader_fn, computed_size, stream_obj = self._prepare_async_reader(stream, size) + + async def _run(): + try: + await self.initiate( + transport=sess, + request_body=request_body, + size=computed_size, + progress_queue=progress_queue, + ) + + final_resp_tuple = None + while not self._state.finished and not self._state.invalid: + final_resp_tuple = await self._transmit_chunk( + sess, reader_fn, computed_size, progress_queue, stream_obj + ) + + if final_resp_tuple is None: + raise ValueError( + "Upload completed without receiving a final response." + ) + + _, _, body_bytes = final_resp_tuple + self._response = self._format_response(body_bytes) + progress_queue.put_nowait(_DONE_SENTINEL) + return self._response + except Exception as exc: + self._enrich_exception(exc) + progress_queue.put_nowait(exc) + raise + + task = asyncio.create_task(_run()) + return AsyncUploadOperation( + task=task, session=self, progress_queue=progress_queue + ) + + 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, + ) -> AsyncUploadOperation: + """Resumes an existing upload asynchronously, returning an AsyncUploadOperation. + + Args: + upload_url: Established upload session URL. + stream: Data payload to resume uploading. + size: Total stream size in bytes, if known. + chunk_size: Optional chunk size override in bytes. + transport: Optional aiohttp client session. + + Returns: + An AsyncUploadOperation handle representing the resumed 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 chunk_size is not None: + self._state._chunk_size = chunk_size + + self._state._resumable_url = upload_url + progress_queue: asyncio.Queue = asyncio.Queue() + reader_fn, computed_size, stream_obj = self._prepare_async_reader(stream, size) + + async def _run(): + try: + await self._recover(sess, stream_obj) + self._notify_progress( + common.ProgressState.OFFSET_RECEIVED, progress_queue + ) + + final_resp_tuple = None + while not self._state.finished and not self._state.invalid: + final_resp_tuple = await self._transmit_chunk( + sess, reader_fn, computed_size, progress_queue, stream_obj + ) + + if final_resp_tuple is None: + raise ValueError( + "Upload resumed but completed without receiving a final response." + ) + + _, _, body_bytes = final_resp_tuple + self._response = self._format_response(body_bytes) + progress_queue.put_nowait(_DONE_SENTINEL) + return self._response + except Exception as exc: + self._enrich_exception(exc) + progress_queue.put_nowait(exc) + raise + + task = asyncio.create_task(_run()) + return AsyncUploadOperation( + task=task, session=self, progress_queue=progress_queue + ) + + def _prepare_async_reader( + self, + stream: Union[AsyncIterable[bytes], BinaryIO, bytes, Iterable[bytes]], + size: Optional[int], + ) -> Tuple[Callable[[int], Awaitable[bytes]], Optional[int], Any]: + """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 + + if hasattr(stream, "read") and inspect.iscoroutinefunction(stream.read): + # Native async reader (e.g. asyncio.StreamReader) + async def reader(n: int) -> bytes: + return await stream.read(n) # type: ignore + + return reader, computed_size, stream + + if hasattr(stream, "read"): + # Synchronous binary stream: offload blocking reads to worker thread + sync_stream: Any = stream + if computed_size is None and hasattr(sync_stream, "getbuffer"): + computed_size = sync_stream.getbuffer().nbytes + + if hasattr(sync_stream, "tell"): + try: + self._start_stream_offset = sync_stream.tell() + except (OSError, AttributeError): + self._start_stream_offset = 0 + + async def reader(n: int) -> bytes: + return await asyncio.to_thread(sync_stream.read, n) + + return reader, computed_size, sync_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)}") + + def _format_response(self, response_bytes: bytes) -> Any: + """Formats response bytes into protobuf message type if provided. + + Args: + response_bytes: Raw HTTP response body bytes from final chunk. + + Returns: + Deserialized protobuf message or the raw bytes response. + """ + return _format_response_payload(response_bytes, self._config.response_type) 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..f748cb01e9cb --- /dev/null +++ b/packages/google-api-core/google/api_core/resumable_transfer/upload_state.py @@ -0,0 +1,316 @@ +# 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, + resumable_url: Optional[str] = None, + ) -> None: + """Initializes the protocol state machine. + + Args: + upload_url: The initial endpoint URL for starting the upload. + chunk_size: Desired chunk size in bytes. + resumable_url: Established upload session URL if resuming. + """ + self._initial_url = upload_url or "" + self._chunk_size = chunk_size + self._resumable_url = resumable_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: + """The initial endpoint URL for starting the upload.""" + return self._initial_url + + @property + def resumable_url(self) -> Optional[str]: + """The established upload session URL, or None if not established.""" + return self._resumable_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). + """ + self._total_bytes = size + 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" + ) + + resumable_url = headers_lower.get(common.HEADER_URL.lower()) + if not resumable_url: + self._invalid = True + raise ValueError(f"Server did not return {common.HEADER_URL} header") + + self._resumable_url = resumable_url + granularity = headers_lower.get(common.HEADER_CHUNK_GRANULARITY.lower()) + if granularity: + self._chunk_granularity = int(granularity) + + return self._resumable_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._resumable_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._resumable_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._resumable_url: + raise ValueError("Upload session URL not established.") + + headers = {common.HEADER_COMMAND: common.Command.QUERY.value} + return "POST", self._resumable_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._resumable_url: + raise ValueError("Upload session URL not established.") + + headers = {common.HEADER_COMMAND: common.Command.CANCEL.value} + return "POST", self._resumable_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..606b928d540b --- /dev/null +++ b/packages/google-api-core/tests/asyncio/test_resumable_transfer_async.py @@ -0,0 +1,1239 @@ +# 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 is None + 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, "aiohttp", None) + 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]) + config = ResumableUploadConfig(response_type=DummyResponse) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + 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, response_type=DummyResponse) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + 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, response_type=DummyResponse) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + 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", + config=ResumableUploadConfig(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", + config=ResumableUploadConfig(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", + config=ResumableUploadConfig(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()) + with pytest.raises(TypeError, match="Unsupported stream type"): + session._prepare_async_reader(stream=12345, size=10) # type: ignore + + +# ===================================================================== +# 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( + config=ResumableUploadConfig(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" + + +@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( + config=ResumableUploadConfig(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( + config=ResumableUploadConfig(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( + resumable_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", + config=ResumableUploadConfig(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 409 Conflict + chunk_conflict = DummyAsyncResponse(status=409, headers={}, body=b"Conflict") + # 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_conflict, query_resp, chunk_success] + ) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=ResumableUploadConfig(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", + config=ResumableUploadConfig(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, 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") + + +@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", + config=ResumableUploadConfig(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", + config=ResumableUploadConfig(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", + config=ResumableUploadConfig(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_without_dict() -> None: + """Verifies that _enrich_exception handles objects without __dict__.""" + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + ) + exc = Exception() + session._enrich_exception(exc) + + +def test_async_notify_progress_branches() -> None: + """Verifies progress notification callbacks and queues.""" + called = [] + config = ResumableUploadConfig(on_progress=lambda p: called.append(p)) + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + ) + # When upload_url is established on state + session._state._resumable_url = "https://upload.example.com/resumable-async" + q: asyncio.Queue = asyncio.Queue() + session._notify_progress(common.ProgressState.UPLOADING, queue=q) + assert len(called) == 1 + assert q.qsize() == 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 branches in _async_retry.""" + session = AsyncResumableUploadSession( + upload_url="https://api.example.com/start", + ) + + # MissingStatusHeaderError retries and raises on final attempt + attempts = 0 + + async def fail_missing_header(): + nonlocal attempts + attempts += 1 + raise exceptions.MissingStatusHeaderError("missing") + + with pytest.raises(exceptions.MissingStatusHeaderError): + await session._async_retry(fail_missing_header, max_attempts=2) + assert attempts == 2 + + # Non-retryable GoogleAPICallError raises immediately + async def fail_400(): + raise exceptions.from_http_status(400, "Bad Request") + + with pytest.raises(exceptions.BadRequest): + await session._async_retry(fail_400, max_attempts=3) + + # Retryable GoogleAPICallError retries and raises on final attempt + attempts_503 = 0 + + async def fail_503(): + nonlocal attempts_503 + attempts_503 += 1 + raise exceptions.from_http_status(503, "Service Unavailable") + + with pytest.raises(exceptions.ServiceUnavailable): + await session._async_retry(fail_503, max_attempts=2) + assert attempts_503 == 2 + + +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( + resumable_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( + resumable_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._resumable_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._resumable_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._resumable_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) diff --git a/packages/google-api-core/tests/helpers.py b/packages/google-api-core/tests/helpers.py index 86b5d149755f..cd55a32e9da5 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..4f5f9be5944b --- /dev/null +++ b/packages/google-api-core/tests/unit/test_resumable_transfer.py @@ -0,0 +1,1298 @@ +# 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, + Command, + MissingStatusHeaderError, + ProgressState, + ResumableUploadConfig, + ResumableUploadSession, + Status, + TransferStalledError, + UnseekableStreamError, + UploadCancelledError, + UploadProgress, + common, + 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 Command.START == "start" + assert Command.UPLOAD == "upload" + assert Command.FINALIZE == "finalize" + assert Command.QUERY == "query" + assert Command.CANCEL == "cancel" + + assert Status.ACTIVE == "active" + assert Status.FINAL == "final" + assert 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.resumable_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( + resumable_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( + resumable_url="https://upload.example.com/session" + ) + with pytest.raises(MissingStatusHeaderError): + state.process_chunk_response(200, {}, 10) + + +def test_protocol_state_query_and_cancel(): + state = upload_state.ProtocolState( + resumable_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 + + 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.start_timeout is None + assert config.start_retry is None + assert config.headers is None + assert config.deadline is None + + +def test_resumable_upload_config_fallbacks_and_headers(): + retry1 = mock.Mock() + config1 = ResumableUploadConfig( + start_timeout=45.0, + retry=retry1, + headers={"X-Test": "1"}, + ) + assert config1.timeout == 45.0 + assert config1.start_timeout == 45.0 + assert config1.retry is retry1 + assert config1.start_retry is retry1 + assert config1.start_headers == [("X-Test", "1")] + + retry2 = mock.Mock() + config2 = ResumableUploadConfig( + timeout=30.0, + start_retry=retry2, + headers=[("X-Test", "2")], + ) + assert config2.timeout == 30.0 + assert config2.start_timeout == 30.0 + assert config2.retry is retry2 + assert config2.start_retry is retry2 + 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] + + config = ResumableUploadConfig(response_type=DummyResponse) + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=session_transport, + ) + + 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, response_type=DummyResponse) + session = ResumableUploadSession( + upload_url="https://api.example.com/start", + config=config, + transport=session_transport, + ) + + 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] + + config = ResumableUploadConfig(response_type=DummyResponse) + session = ResumableUploadSession(config=config) + + 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] + + config = ResumableUploadConfig(response_type=DummyResponse) + session = ResumableUploadSession(config=config) + + 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", + config=ResumableUploadConfig(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_control_timeout(): + 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, + ) + + 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( + resumable_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", + config=ResumableUploadConfig(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", + config=ResumableUploadConfig(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", + config=ResumableUploadConfig(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", + config=ResumableUploadConfig(response_type=None), + transport=session_transport, + ) + resp = session.upload(stream=b"payload") + assert resp is chunk_resp + + +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.resumable_url is None + 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(): + 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("https://api.example.com/init") + 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.resumable_url is None + 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 = session._get_retry_predicate() + + assert pred(exceptions.DeadlineExceeded("deadline")) is False + assert pred(TransferStalledError("stalled")) is False + assert pred(UploadCancelledError("cancelled")) is False + assert pred(MissingStatusHeaderError("missing")) is True + assert pred(requests.exceptions.ConnectionError("conn")) is True + assert pred(requests.exceptions.ChunkedEncodingError("chunked")) is True + assert pred(exceptions.from_http_status(503, "503")) is True + assert pred(exceptions.from_http_status(400, "400")) is False + assert pred(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.read() == b"hello world" + assert computed_size == 11 + + 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_fallbacks_and_headers(): + cfg1 = ResumableUploadConfig(start_timeout=15.0) + assert cfg1.timeout == 15.0 + + cfg2 = ResumableUploadConfig(timeout=25.0) + assert cfg2.start_timeout == 25.0 + + ret = mock.Mock(spec=google.api_core.retry.Retry) + cfg3 = ResumableUploadConfig(start_retry=ret) + assert cfg3.retry is ret + + cfg4 = ResumableUploadConfig(retry=ret) + assert cfg4.start_retry is ret + + 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( + resumable_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(): + callback_mock = mock.Mock() + config = ResumableUploadConfig(on_progress=callback_mock) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config, + ) + session._state._resumable_url = "https://api.example.com/init" + with session._capture_progress() as captured: + session._notify_progress(common.ProgressState.UPLOADING) + assert len(captured) == 1 + assert captured[0].state == common.ProgressState.UPLOADING + assert callback_mock.called + assert callback_mock.call_args[0][0] is captured[0] + + +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_fallback(): + ret_start = mock.Mock(spec=google.api_core.retry.Retry) + ret_fallback = mock.Mock(spec=google.api_core.retry.Retry) + config = ResumableUploadConfig(start_retry=ret_start, retry=ret_fallback) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config, + ) + assert session._get_retry(is_start=True) is ret_start + assert session._get_retry() is ret_fallback + + +def test_sync_stall_control_with_deadline(): + import time + + # 1. compute_chunk_timeout with fallback timeout and stall control active + config = ResumableUploadConfig( + stall_minimum_rate=1024, + stall_timeout=10.0, + timeout=15.0, + ) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config, + ) + t1 = session._compute_chunk_timeout(512) + assert t1 <= 15.0 + + # 2. compute_chunk_timeout with deadline active + 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 + + # 3. update_stall_control raises DeadlineExceeded + 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, time.monotonic() - 15.0, 15.0) + + # 4. update_stall_control raises TransferStalledError (no deadline) + config4 = ResumableUploadConfig( + stall_minimum_rate=1024, + stall_timeout=10.0, + ) + session4 = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config4, + ) + session4._state._resumable_url = "https://api.example.com/init" + with pytest.raises(exceptions.TransferStalledError): + session4._update_stall_control(512, time.monotonic() - 15.0, 15.0) + + +def test_sync_update_stall_control_disabled(): + import time + + 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, time.monotonic(), 5.0) + assert session._aggregate_lag == 0.0 + + +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._resumable_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(): + 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._resumable_url = "https://upload.example.com/resumable-123" + + 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._resumable_url = "https://upload.example.com/resumable-123" + session_dl._get_deadline_remaining = mock.Mock(side_effect=[5.0, -1.0]) + with pytest.raises(exceptions.DeadlineExceeded): + session_dl._transmit_chunk(transport, io.BytesIO(b"data"), size=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, + retry=google.api_core.retry.Retry(predicate=lambda e: False), + ) + session = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config, + transport=transport, + ) + session._state._resumable_url = "https://upload.example.com/resumable-123" + with pytest.raises(exceptions.TransferStalledError): + session._transmit_chunk(transport, io.BytesIO(b"data"), size=4) + + # To hit line 554-558 (outer exception handler with elapsed deadline) + config_dl = ResumableUploadConfig( + stall_minimum_rate=0, + deadline=datetime.datetime.now(datetime.timezone.utc) + + datetime.timedelta(seconds=5), + retry=google.api_core.retry.Retry(predicate=lambda e: False), + ) + session_dl = ResumableUploadSession( + upload_url="https://api.example.com/init", + config=config_dl, + transport=transport, + ) + session_dl._state._resumable_url = "https://upload.example.com/resumable-123" + session_dl._get_deadline_remaining = mock.Mock(side_effect=[5.0, -1.0]) + with pytest.raises(exceptions.DeadlineExceeded): + session_dl._transmit_chunk(transport, io.BytesIO(b"data"), size=4) + + +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._resumable_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 isinstance(res, CustomBytesConvertible) + + 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)