diff --git a/google/cloud/storage/fileio.py b/google/cloud/storage/fileio.py index 639e731ba..97d234983 100644 --- a/google/cloud/storage/fileio.py +++ b/google/cloud/storage/fileio.py @@ -406,9 +406,15 @@ def _upload_chunks_from_buffer(self, num_chunks): upload, transport = self._upload_and_transport + # Attach timeout if specified in the keyword arguments. + # Otherwise, the default timeout will be used from the media library. + kwargs = {} + if "timeout" in self._upload_kwargs: + kwargs = {"timeout": self._upload_kwargs.get("timeout")} + # Upload chunks. The SlidingBuffer class will manage seek position. for _ in range(num_chunks): - upload.transmit_next_chunk(transport) + upload.transmit_next_chunk(transport, **kwargs) # Wipe the buffer of chunks uploaded, preserving any remaining data. self._buffer.flush() diff --git a/tests/unit/test_fileio.py b/tests/unit/test_fileio.py index a92df37da..cafc65e49 100644 --- a/tests/unit/test_fileio.py +++ b/tests/unit/test_fileio.py @@ -346,6 +346,7 @@ def test_write(self, mock_warn): blob = mock.Mock() upload = mock.Mock() transport = mock.Mock() + timeout = 600 blob._initiate_resumable_upload.return_value = (upload, transport) @@ -354,7 +355,10 @@ def test_write(self, mock_warn): # arguments are used. # It would be normal to use a context manager here, but not doing so # gives us more control over close() for test purposes. - upload_kwargs = {"if_metageneration_match": 1} + upload_kwargs = { + "if_metageneration_match": 1, + "timeout": timeout, + } chunk_size = 8 # Note: Real upload requires a multiple of 256KiB. writer = self._make_blob_writer( blob, @@ -366,7 +370,7 @@ def test_write(self, mock_warn): # The transmit_next_chunk method must actually consume bytes from the # sliding buffer for the flush() feature to work properly. - upload.transmit_next_chunk.side_effect = lambda _: writer._buffer.read( + upload.transmit_next_chunk.side_effect = lambda _, timeout: writer._buffer.read( chunk_size ) @@ -388,7 +392,7 @@ def test_write(self, mock_warn): retry=None, **upload_kwargs ) - upload.transmit_next_chunk.assert_called_with(transport) + upload.transmit_next_chunk.assert_called_with(transport, timeout=timeout) self.assertEqual(upload.transmit_next_chunk.call_count, 4) # Write another byte, finalize and close.