Skip to content

Commit

Permalink
Fix: Correctly calculate starting offset for retries of ranged reads (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewsg authored Aug 7, 2024
1 parent 5d2091d commit 34302b4
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 2 deletions.
4 changes: 2 additions & 2 deletions google/resumable_media/requests/download.py
Original file line number Diff line number Diff line change
Expand Up @@ -221,7 +221,7 @@ def retriable_request():
# received using a range request, and set object generation query param.
if self._bytes_downloaded > 0:
_download.add_bytes_range(
self._bytes_downloaded, self.end, self._headers
(self.start or 0) + self._bytes_downloaded, self.end, self._headers
)
request_kwargs["headers"] = self._headers

Expand Down Expand Up @@ -426,7 +426,7 @@ def retriable_request():
# received using a range request, and set object generation query param.
if self._bytes_downloaded > 0:
_download.add_bytes_range(
self._bytes_downloaded, self.end, self._headers
(self.start or 0) + self._bytes_downloaded, self.end, self._headers
)
request_kwargs["headers"] = self._headers

Expand Down
72 changes: 72 additions & 0 deletions tests/unit/requests/test_download.py
Original file line number Diff line number Diff line change
Expand Up @@ -455,6 +455,42 @@ def test_consume_w_bytes_downloaded(self):
range_bytes = "bytes={:d}-{:d}".format(offset, end)
assert download._headers["range"] == range_bytes

def test_consume_w_bytes_downloaded_range_read(self):
stream = io.BytesIO()
chunks = (b"up down ", b"charlie ", b"brown")
start = 1024
end = 65536

download = download_mod.Download(
EXAMPLE_URL,
stream=stream,
start=start,
end=end,
headers=None,
checksum="md5",
)
transport = mock.Mock(spec=["request"])
transport.request.return_value = _mock_response(chunks=chunks, headers=None)

assert download._bytes_downloaded == 0

# Mock a retry operation with bytes already downloaded in the stream and checksum stored
offset = 256
download._bytes_downloaded = offset
download._expected_checksum = None
download._checksum_object = _helpers._DoNothingHash()
download.consume(transport)

called_kwargs = {
"data": None,
"headers": download._headers,
"timeout": EXPECTED_TIMEOUT,
"stream": True,
}
transport.request.assert_called_once_with("GET", EXAMPLE_URL, **called_kwargs)
range_bytes = "bytes={:d}-{:d}".format(offset + start, end)
assert download._headers["range"] == range_bytes

def test_consume_gzip_reset_stream_w_bytes_downloaded(self):
stream = io.BytesIO()
chunks = (b"up down ", b"charlie ", b"brown")
Expand Down Expand Up @@ -904,6 +940,42 @@ def test_consume_w_bytes_downloaded(self):
range_bytes = "bytes={:d}-{:d}".format(offset, end)
assert download._headers["range"] == range_bytes

def test_consume_w_bytes_downloaded_range_read(self):
stream = io.BytesIO()
chunks = (b"up down ", b"charlie ", b"brown")
start = 1024
end = 65536

download = download_mod.RawDownload(
EXAMPLE_URL,
stream=stream,
start=start,
end=end,
headers=None,
checksum="md5",
)
transport = mock.Mock(spec=["request"])
transport.request.return_value = _mock_raw_response(chunks=chunks, headers=None)

assert download._bytes_downloaded == 0

# Mock a retry operation with bytes already downloaded in the stream and checksum stored
offset = 256
download._bytes_downloaded = offset
download._expected_checksum = None
download._checksum_object = _helpers._DoNothingHash()
download.consume(transport)

called_kwargs = {
"data": None,
"headers": download._headers,
"timeout": EXPECTED_TIMEOUT,
"stream": True,
}
transport.request.assert_called_once_with("GET", EXAMPLE_URL, **called_kwargs)
range_bytes = "bytes={:d}-{:d}".format(start + offset, end)
assert download._headers["range"] == range_bytes

def test_consume_gzip_reset_stream_w_bytes_downloaded(self):
stream = io.BytesIO()
chunks = (b"up down ", b"charlie ", b"brown")
Expand Down

0 comments on commit 34302b4

Please sign in to comment.