diff --git a/google/resumable_media/_download.py b/google/resumable_media/_download.py index 7f489d2e..7b67ad20 100644 --- a/google/resumable_media/_download.py +++ b/google/resumable_media/_download.py @@ -41,6 +41,7 @@ class DownloadBase(object): the downloaded resource can be written to. start (int): The first byte in a range to be downloaded. end (int): The last byte in a range to be downloaded. + user_agent (Optional[str]) : Custom user agent to be sent headers (Optional[Mapping[str, str]]): Extra headers that should be sent with the request, e.g. headers for encrypted data. @@ -51,11 +52,12 @@ class DownloadBase(object): """ def __init__(self, media_url, stream=None, - start=None, end=None, headers=None): + start=None, end=None, user_agent=None, headers=None): self.media_url = media_url self._stream = stream self.start = start self.end = end + self.user_agent = user_agent if headers is None: headers = {} self._headers = headers diff --git a/google/resumable_media/requests/download.py b/google/resumable_media/requests/download.py index d0bcd990..9f30cb15 100644 --- a/google/resumable_media/requests/download.py +++ b/google/resumable_media/requests/download.py @@ -64,6 +64,7 @@ class Download(_helpers.RequestsMixin, _download.Download): end (int): The last byte in a range to be downloaded. If not provided, but ``start`` is provided, will download from the ``start`` to the end of the media. + user_agent (Optional[str]) : Custom user agent to be sent headers (Optional[Mapping[str, str]]): Extra headers that should be sent with the request, e.g. headers for encrypted data. @@ -156,6 +157,8 @@ def consume(self, transport): finished. """ method, url, payload, headers = self._prepare_request() + if self.user_agent is not None: + headers[u'User-Agent'] = self.user_agent # NOTE: We assume "payload is None" but pass it along anyway. request_kwargs = { u'data': payload, diff --git a/tests/unit/requests/test_download.py b/tests/unit/requests/test_download.py index 4e1715cd..3c0faca6 100644 --- a/tests/unit/requests/test_download.py +++ b/tests/unit/requests/test_download.py @@ -132,10 +132,11 @@ def test__write_to_stream_with_hash_check_fail(self): decode_unicode=False) def _consume_helper( - self, stream=None, end=65536, headers=None, chunks=(), - response_headers=None): + self, stream=None, end=65536, user_agent=None, headers=None, + chunks=(), response_headers=None): download = download_mod.Download( - EXAMPLE_URL, stream=stream, end=end, headers=headers) + EXAMPLE_URL, stream=stream, end=end, user_agent=user_agent, + headers=headers) transport = mock.Mock(spec=[u'request']) transport.request.return_value = _mock_response( chunks=chunks, headers=response_headers) @@ -234,6 +235,14 @@ def test_consume_with_headers(self): # Make sure the headers have been modified. assert headers == {u'range': range_bytes} + def test_consume_with_user_agent(self): + headers = {} + end = 16383 + user_agent = "Custom-User-Agent-1.0" + range_bytes = u'bytes={:d}-{:d}'.format(0, end) + self._consume_helper(end=end, user_agent=user_agent, headers=headers) + assert headers == {u'range': range_bytes, u'User-Agent': user_agent} + class TestChunkedDownload(object): diff --git a/tests/unit/test__download.py b/tests/unit/test__download.py index ee4e9fc7..a693e353 100644 --- a/tests/unit/test__download.py +++ b/tests/unit/test__download.py @@ -35,6 +35,7 @@ def test_constructor_defaults(self): assert download._stream is None assert download.start is None assert download.end is None + assert download.user_agent is None assert download._headers == {} assert not download._finished _check_retry_strategy(download) @@ -42,14 +43,16 @@ def test_constructor_defaults(self): def test_constructor_explicit(self): start = 11 end = 10001 + user_agent = 'custome-user-agent-1.0' headers = {u'foof': u'barf'} download = _download.DownloadBase( EXAMPLE_URL, stream=mock.sentinel.stream, - start=start, end=end, headers=headers) + start=start, end=end, user_agent=user_agent, headers=headers) assert download.media_url == EXAMPLE_URL assert download._stream is mock.sentinel.stream assert download.start == start assert download.end == end + assert download.user_agent == user_agent assert download._headers is headers assert not download._finished _check_retry_strategy(download)