feat: add context manager to mrd (#1724)

feat: add context manager to mrd

Author: chandra-siri

google/cloud/storage/_experimental/asyncio/async_multi_range_downloader.py

@@ -209,6 +209,16 @@ def __init__(
 self._download_ranges_id_to_pending_read_ids= {}
 self.persisted_size: Optional[int] =None# updated after opening the stream
 
+asyncdef__aenter__(self):
+"""Opens the underlying bidi-gRPC connection to read from the object."""
+awaitself.open()
+returnself
+
+asyncdef__aexit__(self, exc_type, exc_val, exc_tb):
+"""Closes the underlying bidi-gRPC connection."""
+ifself.is_stream_open:
+awaitself.close()
+
 def_on_open_error(self, exc):
 """Extracts routing token and read handle on redirect error during open."""
 routing_token, read_handle=_handle_redirect(exc)

tests/system/test_zonal.py

@@ -116,19 +116,19 @@ async def _run():
 assertobject_metadata.size==object_size
 assertint(object_metadata.checksums.crc32c) ==object_checksum
 
-mrd=AsyncMultiRangeDownloader(grpc_client, _ZONAL_BUCKET, object_name)
 buffer=BytesIO()
-awaitmrd.open()
-# (0, 0) means read the whole object
-awaitmrd.download_ranges([(0, 0, buffer)])
-awaitmrd.close()
+asyncwithAsyncMultiRangeDownloader(
+grpc_client, _ZONAL_BUCKET, object_name
+ ) asmrd:
+# (0, 0) means read the whole object
+awaitmrd.download_ranges([(0, 0, buffer)])
+assertmrd.persisted_size==object_size
+
 assertbuffer.getvalue() ==object_data
-assertmrd.persisted_size==object_size
 
 # Clean up; use json client (i.e. `storage_client` fixture) to delete.
 blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))
 delwriter
-delmrd
 gc.collect()
 
 event_loop.run_until_complete(_run())

tests/unit/asyncio/test_async_multi_range_downloader.py

@@ -401,3 +401,44 @@ async def test_download_ranges_raises_on_checksum_mismatch(
 
 assert"Checksum mismatch"instr(exc_info.value)
 mock_checksum_class.assert_called_once_with(test_data)
+
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_multi_range_downloader.AsyncMultiRangeDownloader.open",
+new_callable=AsyncMock,
+ )
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_multi_range_downloader.AsyncMultiRangeDownloader.close",
+new_callable=AsyncMock,
+ )
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client"
+ )
+@pytest.mark.asyncio
+asyncdeftest_async_context_manager_calls_open_and_close(
+self, mock_grpc_client, mock_close, mock_open
+ ):
+# Arrange
+mrd=AsyncMultiRangeDownloader(
+mock_grpc_client, _TEST_BUCKET_NAME, _TEST_OBJECT_NAME
+ )
+
+# To simulate the behavior of open and close changing the stream state
+asyncdefopen_side_effect():
+mrd._is_stream_open=True
+
+asyncdefclose_side_effect():
+mrd._is_stream_open=False
+
+mock_open.side_effect=open_side_effect
+mock_close.side_effect=close_side_effect
+mrd._is_stream_open=False
+
+# Act
+asyncwithmrdasdownloader:
+# Assert
+mock_open.assert_called_once()
+assertdownloader==mrd
+assertmrd.is_stream_open
+
+mock_close.assert_called_once()
+assertnotmrd.is_stream_open