test: Add tests for checksum validation with partial responses

googleapis · Oct 24, 2022 · 691b93c · 691b93c
1 parent 79bf803
commit 691b93c
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 4 deletions.
diff --git a/google/_async_resumable_media/requests/download.py b/google/_async_resumable_media/requests/download.py
@@ -92,7 +92,10 @@ async def _write_to_stream(self, response):
             local_checksum_object.update(chunk)
 
         # Don't validate the checksum for partial responses.
-        if expected_checksum is not None and response.status != http.client.PARTIAL_CONTENT:
+        if (
+            expected_checksum is not None
+            and response.status != http.client.PARTIAL_CONTENT
+        ):
             actual_checksum = sync_helpers.prepare_checksum_digest(
                 checksum_object.digest()
             )
@@ -216,7 +219,10 @@ async def _write_to_stream(self, response):
             checksum_object.update(chunk)
 
         # Don't validate the checksum for partial responses.
-        if expected_checksum is not None and response.status != http.client.PARTIAL_CONTENT:
+        if (
+            expected_checksum is not None
+            and response.status != http.client.PARTIAL_CONTENT
+        ):
             actual_checksum = sync_helpers.prepare_checksum_digest(
                 checksum_object.digest()
             )

diff --git a/google/resumable_media/requests/download.py b/google/resumable_media/requests/download.py
@@ -126,7 +126,10 @@ def _write_to_stream(self, response):
                 local_checksum_object.update(chunk)
 
         # Don't validate the checksum for partial responses.
-        if expected_checksum is not None and response.status_code != http.client.PARTIAL_CONTENT:
+        if (
+            expected_checksum is not None
+            and response.status_code != http.client.PARTIAL_CONTENT
+        ):
             actual_checksum = _helpers.prepare_checksum_digest(checksum_object.digest())
             if actual_checksum != expected_checksum:
                 msg = _CHECKSUM_MISMATCH.format(
@@ -311,7 +314,10 @@ def _write_to_stream(self, response):
             response._content_consumed = True
 
         # Don't validate the checksum for partial responses.
-        if expected_checksum is not None and response.status_code != http.client.PARTIAL_CONTENT:
+        if (
+            expected_checksum is not None
+            and response.status_code != http.client.PARTIAL_CONTENT
+        ):
             actual_checksum = _helpers.prepare_checksum_digest(checksum_object.digest())
 
             if actual_checksum != expected_checksum:

diff --git a/tests/unit/requests/test_download.py b/tests/unit/requests/test_download.py
@@ -114,6 +114,31 @@ def test__write_to_stream_with_hash_check_fail(self, checksum):
             chunk_size=_request_helpers._SINGLE_GET_CHUNK_SIZE, decode_unicode=False
         )
 
+    @pytest.mark.parametrize("checksum", ["md5", "crc32c"])
+    def test__write_to_stream_no_checksum_validation_for_partial_response(self, checksum):
+        stream = io.BytesIO()
+        download = download_mod.Download(EXAMPLE_URL, stream=stream, checksum=checksum)
+
+        chunk1 = b"first chunk"
+        response = _mock_response(status_code=http.client.PARTIAL_CONTENT, chunks=[chunk1])
+
+        # Make sure that the checksum is not validated.
+        with mock.patch(
+            "google.resumable_media._helpers.prepare_checksum_digest",
+            return_value=None,
+        ) as prepare_checksum_digest:
+            ret_val = download._write_to_stream(response)
+            assert not prepare_checksum_digest.called
+
+        assert not download.finished
+
+        # Check mocks.
+        response.__enter__.assert_called_once_with()
+        response.__exit__.assert_called_once_with(None, None, None)
+        response.iter_content.assert_called_once_with(
+            chunk_size=_request_helpers._SINGLE_GET_CHUNK_SIZE, decode_unicode=False
+        )
+
     def test__write_to_stream_with_invalid_checksum_type(self):
         BAD_CHECKSUM_TYPE = "badsum"
 

diff --git a/tests_async/unit/requests/test_download.py b/tests_async/unit/requests/test_download.py
@@ -100,6 +100,25 @@ async def test__write_to_stream_with_hash_check_fail(self, checksum):
         )
         assert error.args[0] == msg
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("checksum", ["md5", "crc32c"])
+    async def test__write_to_stream_no_checksum_validation_for_partial_response(self, checksum):
+        stream = io.BytesIO()
+        download = download_mod.Download(sync_test.EXAMPLE_URL, stream=stream, checksum=checksum)
+
+        chunk1 = b"first chunk"
+        response = _mock_response(status=http.client.PARTIAL_CONTENT, chunks=[chunk1])
+
+        # Make sure that the checksum is not validated.
+        with mock.patch(
+            "google.resumable_media._helpers.prepare_checksum_digest",
+            return_value=None,
+        ) as prepare_checksum_digest:
+            ret_val = await download._write_to_stream(response)
+            assert not prepare_checksum_digest.called
+
+        assert not download.finished
+
     @pytest.mark.asyncio
     async def test__write_to_stream_with_invalid_checksum_type(self):
         BAD_CHECKSUM_TYPE = "badsum"