diff --git a/google/cloud/storage/_media/_download.py b/google/cloud/storage/_media/_download.py index 349ddf30c..422b98041 100644 --- a/google/cloud/storage/_media/_download.py +++ b/google/cloud/storage/_media/_download.py @@ -140,7 +140,7 @@ class Download(DownloadBase): ``start`` to the end of the media. headers (Optional[Mapping[str, str]]): Extra headers that should be sent with the request, e.g. headers for encrypted data. - checksum Optional([str]): The type of checksum to compute to verify + checksum (Optional[str]): The type of checksum to compute to verify the integrity of the object. The response headers must contain a checksum of the requested type. If the headers lack an appropriate checksum (for instance in the case of transcoded or @@ -157,6 +157,9 @@ class Download(DownloadBase): See the retry.py source code and docstrings in this package (google.cloud.storage.retry) for information on retry types and how to configure them. + single_shot_download (Optional[bool]): If true, download the object in a single request. + Caution: Enabling this will increase the memory overload for your application. + Please enable this as per your use case. """ @@ -169,6 +172,7 @@ def __init__( headers=None, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ): super(Download, self).__init__( media_url, stream=stream, start=start, end=end, headers=headers, retry=retry @@ -178,6 +182,7 @@ def __init__( self.checksum = ( "crc32c" if _helpers._is_crc32c_available_and_fast() else "md5" ) + self.single_shot_download = single_shot_download self._bytes_downloaded = 0 self._expected_checksum = None self._checksum_object = None diff --git a/google/cloud/storage/_media/requests/download.py b/google/cloud/storage/_media/requests/download.py index 67535f923..b8e2758e1 100644 --- a/google/cloud/storage/_media/requests/download.py +++ b/google/cloud/storage/_media/requests/download.py @@ -132,13 +132,24 @@ def _write_to_stream(self, response): # the stream is indeed compressed, this will delegate the checksum # object to the decoder and return a _DoNothingHash here. local_checksum_object = _add_decoder(response.raw, checksum_object) - body_iter = response.iter_content( - chunk_size=_request_helpers._SINGLE_GET_CHUNK_SIZE, decode_unicode=False - ) - for chunk in body_iter: - self._stream.write(chunk) - self._bytes_downloaded += len(chunk) - local_checksum_object.update(chunk) + + # This is useful for smaller files, or when the user wants to + # download the entire file in one go. + if self.single_shot_download: + content = response.raw.read(decode_content=True) + self._stream.write(content) + self._bytes_downloaded += len(content) + local_checksum_object.update(content) + response._content_consumed = True + else: + body_iter = response.iter_content( + chunk_size=_request_helpers._SINGLE_GET_CHUNK_SIZE, + decode_unicode=False, + ) + for chunk in body_iter: + self._stream.write(chunk) + self._bytes_downloaded += len(chunk) + local_checksum_object.update(chunk) # Don't validate the checksum for partial responses. if ( @@ -345,13 +356,21 @@ def _write_to_stream(self, response): checksum_object = self._checksum_object with response: - body_iter = response.raw.stream( - _request_helpers._SINGLE_GET_CHUNK_SIZE, decode_content=False - ) - for chunk in body_iter: - self._stream.write(chunk) - self._bytes_downloaded += len(chunk) - checksum_object.update(chunk) + # This is useful for smaller files, or when the user wants to + # download the entire file in one go. + if self.single_shot_download: + content = response.raw.read() + self._stream.write(content) + self._bytes_downloaded += len(content) + checksum_object.update(content) + else: + body_iter = response.raw.stream( + _request_helpers._SINGLE_GET_CHUNK_SIZE, decode_content=False + ) + for chunk in body_iter: + self._stream.write(chunk) + self._bytes_downloaded += len(chunk) + checksum_object.update(chunk) response._content_consumed = True # Don't validate the checksum for partial responses. diff --git a/google/cloud/storage/blob.py b/google/cloud/storage/blob.py index 0d0e8ee80..0eb94fd47 100644 --- a/google/cloud/storage/blob.py +++ b/google/cloud/storage/blob.py @@ -987,6 +987,7 @@ def _do_download( timeout=_DEFAULT_TIMEOUT, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ): """Perform a download without any error handling. @@ -1047,6 +1048,12 @@ def _do_download( See the retry.py source code and docstrings in this package (google.cloud.storage.retry) for information on retry types and how to configure them. + + :type single_shot_download: bool + :param single_shot_download: + (Optional) If true, download the object in a single request. + Caution: Enabling this will increase the memory overload for your application. + Please enable this as per your use case. """ extra_attributes = { @@ -1054,6 +1061,7 @@ def _do_download( "download.chunk_size": f"{self.chunk_size}", "download.raw_download": raw_download, "upload.checksum": f"{checksum}", + "download.single_shot_download": single_shot_download, } args = {"timeout": timeout} @@ -1073,6 +1081,10 @@ def _do_download( end=end, checksum=checksum, retry=retry, + # NOTE: single_shot_download is only supported in Download and RawDownload + # classes, i.e., when chunk_size is set to None (the default value). It is + # not supported for chunked downloads. + single_shot_download=single_shot_download, ) with create_trace_span( name=f"Storage.{download_class}/consume", @@ -1127,6 +1139,7 @@ def download_to_file( timeout=_DEFAULT_TIMEOUT, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ): """Download the contents of this blob into a file-like object. @@ -1222,6 +1235,12 @@ def download_to_file( (google.cloud.storage.retry) for information on retry types and how to configure them. + :type single_shot_download: bool + :param single_shot_download: + (Optional) If true, download the object in a single request. + Caution: Enabling this will increase the memory overload for your application. + Please enable this as per your use case. + :raises: :class:`google.cloud.exceptions.NotFound` """ with create_trace_span(name="Storage.Blob.downloadToFile"): @@ -1240,6 +1259,7 @@ def download_to_file( timeout=timeout, checksum=checksum, retry=retry, + single_shot_download=single_shot_download, ) def _handle_filename_and_download(self, filename, *args, **kwargs): @@ -1285,6 +1305,7 @@ def download_to_filename( timeout=_DEFAULT_TIMEOUT, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ): """Download the contents of this blob into a named file. @@ -1370,6 +1391,12 @@ def download_to_filename( (google.cloud.storage.retry) for information on retry types and how to configure them. + :type single_shot_download: bool + :param single_shot_download: + (Optional) If true, download the object in a single request. + Caution: Enabling this will increase the memory overload for your application. + Please enable this as per your use case. + :raises: :class:`google.cloud.exceptions.NotFound` """ with create_trace_span(name="Storage.Blob.downloadToFilename"): @@ -1388,6 +1415,7 @@ def download_to_filename( timeout=timeout, checksum=checksum, retry=retry, + single_shot_download=single_shot_download, ) def download_as_bytes( @@ -1405,6 +1433,7 @@ def download_as_bytes( timeout=_DEFAULT_TIMEOUT, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ): """Download the contents of this blob as a bytes object. @@ -1484,6 +1513,12 @@ def download_as_bytes( (google.cloud.storage.retry) for information on retry types and how to configure them. + :type single_shot_download: bool + :param single_shot_download: + (Optional) If true, download the object in a single request. + Caution: Enabling this will increase the memory overload for your application. + Please enable this as per your use case. + :rtype: bytes :returns: The data stored in this blob. @@ -1507,6 +1542,7 @@ def download_as_bytes( timeout=timeout, checksum=checksum, retry=retry, + single_shot_download=single_shot_download, ) return string_buffer.getvalue() @@ -1524,6 +1560,7 @@ def download_as_string( if_metageneration_not_match=None, timeout=_DEFAULT_TIMEOUT, retry=DEFAULT_RETRY, + single_shot_download=False, ): """(Deprecated) Download the contents of this blob as a bytes object. @@ -1594,6 +1631,12 @@ def download_as_string( (google.cloud.storage.retry) for information on retry types and how to configure them. + :type single_shot_download: bool + :param single_shot_download: + (Optional) If true, download the object in a single request. + Caution: Enabling this will increase the memory overload for your application. + Please enable this as per your use case. + :rtype: bytes :returns: The data stored in this blob. @@ -1616,6 +1659,7 @@ def download_as_string( if_metageneration_not_match=if_metageneration_not_match, timeout=timeout, retry=retry, + single_shot_download=single_shot_download, ) def download_as_text( @@ -1633,6 +1677,7 @@ def download_as_text( if_metageneration_not_match=None, timeout=_DEFAULT_TIMEOUT, retry=DEFAULT_RETRY, + single_shot_download=False, ): """Download the contents of this blob as text (*not* bytes). @@ -1705,6 +1750,12 @@ def download_as_text( (google.cloud.storage.retry) for information on retry types and how to configure them. + :type single_shot_download: bool + :param single_shot_download: + (Optional) If true, download the object in a single request. + Caution: Enabling this will increase the memory overload for your application. + Please enable this as per your use case. + :rtype: text :returns: The data stored in this blob, decoded to text. """ @@ -1722,6 +1773,7 @@ def download_as_text( if_metageneration_not_match=if_metageneration_not_match, timeout=timeout, retry=retry, + single_shot_download=single_shot_download, ) if encoding is not None: @@ -4019,6 +4071,7 @@ def open( For downloads only, the following additional arguments are supported: - ``raw_download`` + - ``single_shot_download`` For uploads only, the following additional arguments are supported: @@ -4209,6 +4262,7 @@ def _prep_and_do_download( timeout=_DEFAULT_TIMEOUT, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, command=None, ): """Download the contents of a blob object into a file-like object. @@ -4294,6 +4348,12 @@ def _prep_and_do_download( (google.cloud.storage.retry) for information on retry types and how to configure them. + :type single_shot_download: bool + :param single_shot_download: + (Optional) If true, download the object in a single request. + Caution: Enabling this will increase the memory overload for your application. + Please enable this as per your use case. + :type command: str :param command: (Optional) Information about which interface for download was used, @@ -4349,6 +4409,7 @@ def _prep_and_do_download( timeout=timeout, checksum=checksum, retry=retry, + single_shot_download=single_shot_download, ) except InvalidResponse as exc: _raise_from_invalid_response(exc) diff --git a/google/cloud/storage/client.py b/google/cloud/storage/client.py index ba94b26fc..2f56d8719 100644 --- a/google/cloud/storage/client.py +++ b/google/cloud/storage/client.py @@ -1143,6 +1143,7 @@ def download_blob_to_file( timeout=_DEFAULT_TIMEOUT, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ): """Download the contents of a blob object or blob URI into a file-like object. @@ -1216,6 +1217,9 @@ def download_blob_to_file( See the retry.py source code and docstrings in this package (google.cloud.storage.retry) for information on retry types and how to configure them. + + single_shot_download (bool): + (Optional) If true, download the object in a single request. """ with create_trace_span(name="Storage.Client.downloadBlobToFile"): if not isinstance(blob_or_uri, Blob): @@ -1236,6 +1240,7 @@ def download_blob_to_file( timeout=timeout, checksum=checksum, retry=retry, + single_shot_download=single_shot_download, ) def list_blobs( diff --git a/google/cloud/storage/fileio.py b/google/cloud/storage/fileio.py index 289a09cee..7c30f39be 100644 --- a/google/cloud/storage/fileio.py +++ b/google/cloud/storage/fileio.py @@ -35,6 +35,7 @@ "timeout", "retry", "raw_download", + "single_shot_download", } # Valid keyword arguments for upload methods. @@ -99,8 +100,9 @@ class BlobReader(io.BufferedIOBase): - ``if_metageneration_not_match`` - ``timeout`` - ``raw_download`` + - ``single_shot_download`` - Note that download_kwargs (excluding ``raw_download``) are also applied to blob.reload(), + Note that download_kwargs (excluding ``raw_download`` and ``single_shot_download``) are also applied to blob.reload(), if a reload is needed during seek(). """ @@ -177,7 +179,9 @@ def seek(self, pos, whence=0): if self._blob.size is None: reload_kwargs = { - k: v for k, v in self._download_kwargs.items() if k != "raw_download" + k: v + for k, v in self._download_kwargs.items() + if (k != "raw_download" and k != "single_shot_download") } self._blob.reload(**reload_kwargs) diff --git a/tests/resumable_media/system/requests/test_download.py b/tests/resumable_media/system/requests/test_download.py index 04c7246f6..84c44c94c 100644 --- a/tests/resumable_media/system/requests/test_download.py +++ b/tests/resumable_media/system/requests/test_download.py @@ -286,6 +286,23 @@ def test_download_full(self, add_files, authorized_transport, checksum): assert self._read_response_content(response) == actual_contents check_tombstoned(download, authorized_transport) + @pytest.mark.parametrize("checksum", ["auto", "md5", "crc32c", None]) + def test_single_shot_download_full(self, add_files, authorized_transport, checksum): + for info in ALL_FILES: + actual_contents = self._get_contents(info) + blob_name = get_blob_name(info) + + # Create the actual download object. + media_url = utils.DOWNLOAD_URL_TEMPLATE.format(blob_name=blob_name) + download = self._make_one( + media_url, checksum=checksum, single_shot_download=True + ) + # Consume the resource with single_shot_download enabled. + response = download.consume(authorized_transport) + assert response.status_code == http.client.OK + assert self._read_response_content(response) == actual_contents + check_tombstoned(download, authorized_transport) + def test_download_to_stream(self, add_files, authorized_transport): for info in ALL_FILES: actual_contents = self._get_contents(info) @@ -306,6 +323,29 @@ def test_download_to_stream(self, add_files, authorized_transport): assert stream.getvalue() == actual_contents check_tombstoned(download, authorized_transport) + @pytest.mark.parametrize("checksum", ["auto", "md5", "crc32c", None]) + def test_single_shot_download_to_stream(self, add_files, authorized_transport, checksum): + for info in ALL_FILES: + actual_contents = self._get_contents(info) + blob_name = get_blob_name(info) + + # Create the actual download object. + media_url = utils.DOWNLOAD_URL_TEMPLATE.format(blob_name=blob_name) + stream = io.BytesIO() + download = self._make_one( + media_url, checksum=checksum, stream=stream, single_shot_download=True + ) + # Consume the resource with single_shot_download enabled. + response = download.consume(authorized_transport) + assert response.status_code == http.client.OK + with pytest.raises(RuntimeError) as exc_info: + getattr(response, "content") + assert exc_info.value.args == (NO_BODY_ERR,) + assert response._content is False + assert response._content_consumed is True + assert stream.getvalue() == actual_contents + check_tombstoned(download, authorized_transport) + def test_download_gzip_w_stored_content_headers( self, add_files, authorized_transport ): diff --git a/tests/resumable_media/unit/requests/test_download.py b/tests/resumable_media/unit/requests/test_download.py index 568d3238c..b17fbb905 100644 --- a/tests/resumable_media/unit/requests/test_download.py +++ b/tests/resumable_media/unit/requests/test_download.py @@ -213,6 +213,25 @@ def test__write_to_stream_incomplete_read(self, checksum): in error.args[0] ) + @pytest.mark.parametrize("checksum", ["auto", "md5", "crc32c", None]) + def test__write_to_stream_single_shot_download(self, checksum): + stream = io.BytesIO() + download = download_mod.Download( + EXAMPLE_URL, stream=stream, checksum=checksum, single_shot_download=True + ) + + chunk1 = b"all at once!" + response = _mock_response(chunks=[chunk1], headers={}) + ret_val = download._write_to_stream(response) + + assert ret_val is None + assert stream.getvalue() == chunk1 + assert download._bytes_downloaded == len(chunk1) + + response.__enter__.assert_called_once_with() + response.__exit__.assert_called_once_with(None, None, None) + response.raw.read.assert_called_once_with(decode_content=True) + def _consume_helper( self, stream=None, @@ -692,6 +711,24 @@ def test__write_to_stream_incomplete_read(self, checksum): in error.args[0] ) + def test__write_to_stream_single_shot_download(self): + stream = io.BytesIO() + download = download_mod.RawDownload( + EXAMPLE_URL, stream=stream, single_shot_download=True + ) + + chunk1 = b"all at once, raw!" + response = _mock_raw_response(chunks=[chunk1], headers={}) + ret_val = download._write_to_stream(response) + + assert ret_val is None + assert stream.getvalue() == chunk1 + assert download._bytes_downloaded == len(chunk1) + + response.__enter__.assert_called_once_with() + response.__exit__.assert_called_once_with(None, None, None) + response.raw.read.assert_called_once_with() + def _consume_helper( self, stream=None, @@ -1333,6 +1370,7 @@ def _mock_response(status_code=http.client.OK, chunks=None, headers=None): response.__enter__.return_value = response response.__exit__.return_value = None response.iter_content.return_value = iter(chunks) + response.raw.read = mock.Mock(side_effect=lambda *args, **kwargs: b"".join(chunks)) return response else: return mock.Mock( @@ -1348,6 +1386,7 @@ def _mock_raw_response(status_code=http.client.OK, chunks=(), headers=None): mock_raw = mock.Mock(headers=headers, spec=["stream"]) mock_raw.stream.return_value = iter(chunks) + mock_raw.read = mock.Mock(return_value=b"".join(chunks)) response = mock.MagicMock( headers=headers, status_code=int(status_code), diff --git a/tests/system/test_blob.py b/tests/system/test_blob.py index 00f218534..8b50322ba 100644 --- a/tests/system/test_blob.py +++ b/tests/system/test_blob.py @@ -1149,3 +1149,22 @@ def test_object_retention_lock(storage_client, buckets_to_delete, blobs_to_delet blob.retention.retain_until_time = None blob.patch(override_unlocked_retention=True) assert blob.retention.mode is None + + +def test_blob_download_as_bytes_single_shot_download( + shared_bucket, blobs_to_delete, file_data, service_account +): + blob_name = f"download-single-shot-{uuid.uuid4().hex}" + info = file_data["simple"] + with open(info["path"], "rb") as f: + payload = f.read() + + blob = shared_bucket.blob(blob_name) + blob.upload_from_string(payload) + blobs_to_delete.append(blob) + + result_regular_download = blob.download_as_bytes(single_shot_download=False) + assert result_regular_download == payload + + result_single_shot_download = blob.download_as_bytes(single_shot_download=True) + assert result_single_shot_download == payload diff --git a/tests/unit/test_blob.py b/tests/unit/test_blob.py index 06ba62220..937bebaf5 100644 --- a/tests/unit/test_blob.py +++ b/tests/unit/test_blob.py @@ -1282,6 +1282,7 @@ def _do_download_helper_wo_chunks( end=3, checksum="auto", retry=retry, + single_shot_download=False, ) else: patched.assert_called_once_with( @@ -1292,6 +1293,7 @@ def _do_download_helper_wo_chunks( end=None, checksum="auto", retry=retry, + single_shot_download=False, ) patched.return_value.consume.assert_called_once_with( @@ -1499,6 +1501,7 @@ def test_download_to_file_with_failure(self): timeout=expected_timeout, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ) def test_download_to_file_wo_media_link(self): @@ -1530,6 +1533,7 @@ def test_download_to_file_wo_media_link(self): timeout=expected_timeout, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ) def test_download_to_file_w_etag_match(self): @@ -1557,6 +1561,7 @@ def test_download_to_file_w_etag_match(self): timeout=expected_timeout, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ) def test_download_to_file_w_generation_match(self): @@ -1584,10 +1589,16 @@ def test_download_to_file_w_generation_match(self): timeout=expected_timeout, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ) def _download_to_file_helper( - self, use_chunks, raw_download, timeout=None, **extra_kwargs + self, + use_chunks, + raw_download, + timeout=None, + single_shot_download=False, + **extra_kwargs, ): blob_name = "blob-name" client = self._make_client() @@ -1612,9 +1623,16 @@ def _download_to_file_helper( with mock.patch.object(blob, "_prep_and_do_download"): if raw_download: - blob.download_to_file(file_obj, raw_download=True, **extra_kwargs) + blob.download_to_file( + file_obj, + raw_download=True, + single_shot_download=single_shot_download, + **extra_kwargs, + ) else: - blob.download_to_file(file_obj, **extra_kwargs) + blob.download_to_file( + file_obj, single_shot_download=single_shot_download, **extra_kwargs + ) expected_retry = extra_kwargs.get("retry", DEFAULT_RETRY) blob._prep_and_do_download.assert_called_once_with( @@ -1632,6 +1650,7 @@ def _download_to_file_helper( timeout=expected_timeout, checksum="auto", retry=expected_retry, + single_shot_download=single_shot_download, ) def test_download_to_file_wo_chunks_wo_raw(self): @@ -1643,6 +1662,26 @@ def test_download_to_file_wo_chunks_no_retry(self): def test_download_to_file_w_chunks_wo_raw(self): self._download_to_file_helper(use_chunks=True, raw_download=False) + def test_download_to_file_wo_single_shot_download_wo_raw(self): + self._download_to_file_helper( + use_chunks=False, raw_download=False, single_shot_download=False + ) + + def test_download_to_file_w_single_shot_download_wo_raw(self): + self._download_to_file_helper( + use_chunks=False, raw_download=False, single_shot_download=True + ) + + def test_download_to_file_wo_single_shot_download_w_raw(self): + self._download_to_file_helper( + use_chunks=False, raw_download=True, single_shot_download=False + ) + + def test_download_to_file_w_single_shot_download_w_raw(self): + self._download_to_file_helper( + use_chunks=False, raw_download=True, single_shot_download=True + ) + def test_download_to_file_wo_chunks_w_raw(self): self._download_to_file_helper(use_chunks=False, raw_download=True) @@ -1711,6 +1750,7 @@ def _download_to_filename_helper( timeout=expected_timeout, checksum="auto", retry=expected_retry, + single_shot_download=False, ) stream = blob._prep_and_do_download.mock_calls[0].args[0] self.assertEqual(stream.name, temp.name) @@ -1767,6 +1807,7 @@ def test_download_to_filename_w_etag_match(self): timeout=expected_timeout, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ) stream = blob._prep_and_do_download.mock_calls[0].args[0] self.assertEqual(stream.name, temp.name) @@ -1800,6 +1841,7 @@ def test_download_to_filename_w_generation_match(self): timeout=expected_timeout, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ) stream = blob._prep_and_do_download.mock_calls[0].args[0] self.assertEqual(stream.name, temp.name) @@ -1842,6 +1884,7 @@ def test_download_to_filename_corrupted(self): timeout=expected_timeout, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ) stream = blob._prep_and_do_download.mock_calls[0].args[0] self.assertEqual(stream.name, filename) @@ -1884,11 +1927,14 @@ def test_download_to_filename_notfound(self): timeout=expected_timeout, checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ) stream = blob._prep_and_do_download.mock_calls[0].args[0] self.assertEqual(stream.name, filename) - def _download_as_bytes_helper(self, raw_download, timeout=None, **extra_kwargs): + def _download_as_bytes_helper( + self, raw_download, timeout=None, single_shot_download=False, **extra_kwargs + ): blob_name = "blob-name" client = self._make_client() bucket = _Bucket(client) @@ -1898,12 +1944,17 @@ def _download_as_bytes_helper(self, raw_download, timeout=None, **extra_kwargs): if timeout is None: expected_timeout = self._get_default_timeout() fetched = blob.download_as_bytes( - raw_download=raw_download, **extra_kwargs + raw_download=raw_download, + single_shot_download=single_shot_download, + **extra_kwargs, ) else: expected_timeout = timeout fetched = blob.download_as_bytes( - raw_download=raw_download, timeout=timeout, **extra_kwargs + raw_download=raw_download, + timeout=timeout, + single_shot_download=single_shot_download, + **extra_kwargs, ) self.assertEqual(fetched, b"") @@ -1924,6 +1975,7 @@ def _download_as_bytes_helper(self, raw_download, timeout=None, **extra_kwargs): timeout=expected_timeout, checksum="auto", retry=expected_retry, + single_shot_download=single_shot_download, ) stream = blob._prep_and_do_download.mock_calls[0].args[0] self.assertIsInstance(stream, io.BytesIO) @@ -1959,6 +2011,7 @@ def test_download_as_bytes_w_etag_match(self): timeout=self._get_default_timeout(), checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ) def test_download_as_bytes_w_generation_match(self): @@ -1989,6 +2042,7 @@ def test_download_as_bytes_w_generation_match(self): timeout=self._get_default_timeout(), checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ) def test_download_as_bytes_wo_raw(self): @@ -2003,6 +2057,16 @@ def test_download_as_bytes_w_raw(self): def test_download_as_byte_w_custom_timeout(self): self._download_as_bytes_helper(raw_download=False, timeout=9.58) + def test_download_as_bytes_wo_single_shot_download(self): + self._download_as_bytes_helper( + raw_download=False, retry=None, single_shot_download=False + ) + + def test_download_as_bytes_w_single_shot_download(self): + self._download_as_bytes_helper( + raw_download=False, retry=None, single_shot_download=True + ) + def _download_as_text_helper( self, raw_download, @@ -2100,6 +2164,7 @@ def _download_as_text_helper( if_metageneration_match=if_metageneration_match, if_metageneration_not_match=if_metageneration_not_match, retry=expected_retry, + single_shot_download=False, ) def test_download_as_text_wo_raw(self): @@ -2226,6 +2291,7 @@ def test_download_as_string(self, mock_warn): timeout=self._get_default_timeout(), checksum="auto", retry=DEFAULT_RETRY, + single_shot_download=False, ) mock_warn.assert_any_call( @@ -2264,6 +2330,7 @@ def test_download_as_string_no_retry(self, mock_warn): timeout=self._get_default_timeout(), checksum="auto", retry=None, + single_shot_download=False, ) mock_warn.assert_any_call( diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index b671cc092..db8094a95 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -1866,6 +1866,7 @@ def test_download_blob_to_file_with_failure(self): checksum="auto", timeout=_DEFAULT_TIMEOUT, retry=DEFAULT_RETRY, + single_shot_download=False, ) def test_download_blob_to_file_with_uri(self): @@ -1905,6 +1906,7 @@ def test_download_blob_to_file_with_uri(self): checksum="auto", timeout=_DEFAULT_TIMEOUT, retry=DEFAULT_RETRY, + single_shot_download=False, ) def test_download_blob_to_file_with_invalid_uri(self): @@ -2032,6 +2034,7 @@ def _download_blob_to_file_helper( checksum="auto", timeout=_DEFAULT_TIMEOUT, retry=expected_retry, + single_shot_download=False, ) def test_download_blob_to_file_wo_chunks_wo_raw(self):