Skip to content

Commit

Permalink
Adding error remapping for Blob.download_to_file(). (googleapis#3338)
Browse files Browse the repository at this point in the history
  • Loading branch information
dhermes committed Apr 27, 2017
1 parent eaf7bd4 commit 3501c26
Show file tree
Hide file tree
Showing 2 changed files with 130 additions and 20 deletions.
47 changes: 37 additions & 10 deletions storage/google/cloud/storage/blob.py
Original file line number Diff line number Diff line change
Expand Up @@ -338,6 +338,36 @@ def _get_download_url(self):
else:
return self.media_link

def _do_download(self, transport, file_obj, download_url, headers):
"""Perform a download without any error handling.
This is intended to be called by :meth:`download_to_file` so it can
be wrapped with error handling / remapping.
:type transport:
:class:`~google.auth.transport.requests.AuthorizedSession`
:param transport: The transport (with credentials) that will
make authenticated requests.
:type file_obj: file
:param file_obj: A file handle to which to write the blob's data.
:type download_url: str
:param download_url: The URL where the media can be accessed.
:type headers: dict
:param headers: Optional headers to be sent with the request(s).
"""
if self.chunk_size is None:
download = resumable_media.Download(download_url, headers=headers)
response = download.consume(transport)
file_obj.write(response.content)
else:
download = resumable_media.ChunkedDownload(
download_url, self.chunk_size, file_obj, headers=headers)
while not download.finished:
download.consume_next_chunk(transport)

def download_to_file(self, file_obj, client=None):
"""Download the contents of this blob into a file-like object.
Expand Down Expand Up @@ -378,16 +408,13 @@ def download_to_file(self, file_obj, client=None):
transport = google.auth.transport.requests.AuthorizedSession(
client._credentials)

# Download the content.
if self.chunk_size is None:
download = resumable_media.Download(download_url, headers=headers)
response = download.consume(transport)
file_obj.write(response.content)
else:
download = resumable_media.ChunkedDownload(
download_url, self.chunk_size, file_obj, headers=headers)
while not download.finished:
download.consume_next_chunk(transport)
try:
self._do_download(transport, file_obj, download_url, headers)
except resumable_media.InvalidResponse as exc:
response = exc.response
faux_response = httplib2.Response({'status': response.status_code})
raise make_exception(faux_response, response.content,
error_info=download_url, use_json=False)

def download_to_filename(self, filename, client=None):
"""Download the contents of this blob into a named file.
Expand Down
103 changes: 93 additions & 10 deletions storage/tests/unit/test_blob.py
Original file line number Diff line number Diff line change
Expand Up @@ -381,11 +381,102 @@ def _check_session_mocks(self, client, fake_session_factory,
'GET', expected_url, data=None, headers=headers)
self.assertEqual(fake_transport.request.mock_calls, [call, call])

def test__do_download_simple(self):
from io import BytesIO
from six.moves import http_client

blob_name = 'blob-name'
# Create a fake client/bucket and use them in the Blob() constructor.
client = mock.Mock(
_credentials=_make_credentials(), spec=['_credentials'])
bucket = _Bucket(client)
blob = self._make_one(blob_name, bucket=bucket)

# Make sure this will not be chunked.
self.assertIsNone(blob.chunk_size)

transport = mock.Mock(spec=['request'])
transport.request.return_value = self._mock_requests_response(
http_client.OK,
{'content-length': '6', 'content-range': 'bytes 0-5/6'},
content=b'abcdef')
file_obj = BytesIO()
download_url = 'http://test.invalid'
headers = {}
blob._do_download(transport, file_obj, download_url, headers)
# Make sure the download was as expected.
self.assertEqual(file_obj.getvalue(), b'abcdef')

transport.request.assert_called_once_with(
'GET', download_url, data=None, headers=headers)

def test__do_download_chunked(self):
from io import BytesIO

blob_name = 'blob-name'
# Create a fake client/bucket and use them in the Blob() constructor.
client = mock.Mock(
_credentials=_make_credentials(), spec=['_credentials'])
bucket = _Bucket(client)
blob = self._make_one(blob_name, bucket=bucket)

# Modify the blob so there there will be 2 chunks of size 3.
blob._CHUNK_SIZE_MULTIPLE = 1
blob.chunk_size = 3

transport = self._mock_transport()
file_obj = BytesIO()
download_url = 'http://test.invalid'
headers = {}
blob._do_download(transport, file_obj, download_url, headers)
# Make sure the download was as expected.
self.assertEqual(file_obj.getvalue(), b'abcdef')

# Check that the transport was called exactly twice.
self.assertEqual(transport.request.call_count, 2)
# ``headers`` was modified (in place) once for each API call.
self.assertEqual(headers, {'range': 'bytes=3-5'})
call = mock.call(
'GET', download_url, data=None, headers=headers)
self.assertEqual(transport.request.mock_calls, [call, call])

@mock.patch('google.auth.transport.requests.AuthorizedSession')
def test_download_to_file_with_failure(self, fake_session_factory):
from io import BytesIO
from six.moves import http_client
from google.cloud import exceptions

blob_name = 'blob-name'
transport = mock.Mock(spec=['request'])
bad_response_headers = {
'Content-Length': '9',
'Content-Type': 'text/html; charset=UTF-8',
}
transport.request.return_value = self._mock_requests_response(
http_client.NOT_FOUND, bad_response_headers, content=b'Not found')
fake_session_factory.return_value = transport
# Create a fake client/bucket and use them in the Blob() constructor.
client = mock.Mock(
_credentials=_make_credentials(), spec=['_credentials'])
bucket = _Bucket(client)
blob = self._make_one(blob_name, bucket=bucket)
# Set the media link on the blob
blob._properties['mediaLink'] = 'http://test.invalid'

file_obj = BytesIO()
with self.assertRaises(exceptions.NotFound):
blob.download_to_file(file_obj)

self.assertEqual(file_obj.tell(), 0)
# Check that exactly one transport was created.
fake_session_factory.assert_called_once_with(client._credentials)
# Check that the transport was called once.
transport.request.assert_called_once_with(
'GET', blob.media_link, data=None, headers={})

@mock.patch('google.auth.transport.requests.AuthorizedSession')
def test_download_to_file_wo_media_link(self, fake_session_factory):
from io import BytesIO
from six.moves.http_client import OK
from six.moves.http_client import PARTIAL_CONTENT

blob_name = 'blob-name'
fake_session_factory.return_value = self._mock_transport()
Expand Down Expand Up @@ -413,7 +504,6 @@ def test_download_to_file_wo_media_link(self, fake_session_factory):
def _download_to_file_helper(self, fake_session_factory, use_chunks=False):
from io import BytesIO
from six.moves.http_client import OK
from six.moves.http_client import PARTIAL_CONTENT

blob_name = 'blob-name'
fake_transport = self._mock_transport()
Expand Down Expand Up @@ -459,8 +549,6 @@ def test_download_to_file_with_chunk_size(self):
def test_download_to_filename(self, fake_session_factory):
import os
import time
from six.moves.http_client import OK
from six.moves.http_client import PARTIAL_CONTENT
from google.cloud._testing import _NamedTemporaryFile

blob_name = 'blob-name'
Expand Down Expand Up @@ -493,8 +581,6 @@ def test_download_to_filename(self, fake_session_factory):
def test_download_to_filename_w_key(self, fake_session_factory):
import os
import time
from six.moves.http_client import OK
from six.moves.http_client import PARTIAL_CONTENT
from google.cloud._testing import _NamedTemporaryFile

blob_name = 'blob-name'
Expand Down Expand Up @@ -535,9 +621,6 @@ def test_download_to_filename_w_key(self, fake_session_factory):

@mock.patch('google.auth.transport.requests.AuthorizedSession')
def test_download_as_string(self, fake_session_factory):
from six.moves.http_client import OK
from six.moves.http_client import PARTIAL_CONTENT

blob_name = 'blob-name'
fake_session_factory.return_value = self._mock_transport()
# Create a fake client/bucket and use them in the Blob() constructor.
Expand Down

0 comments on commit 3501c26

Please sign in to comment.