diff --git a/google/cloud/storage/_experimental/asyncio/async_multi_range_downloader.py b/google/cloud/storage/_experimental/asyncio/async_multi_range_downloader.py index 340f087da..facad4929 100644 --- a/google/cloud/storage/_experimental/asyncio/async_multi_range_downloader.py +++ b/google/cloud/storage/_experimental/asyncio/async_multi_range_downloader.py @@ -128,10 +128,11 @@ async def create_mrd( client: AsyncGrpcClient, bucket_name: str, object_name: str, - generation_number: Optional[int] = None, + generation: Optional[int] = None, read_handle: Optional[_storage_v2.BidiReadHandle] = None, retry_policy: Optional[AsyncRetry] = None, metadata: Optional[List[Tuple[str, str]]] = None, + **kwargs, ) -> AsyncMultiRangeDownloader: """Initializes a MultiRangeDownloader and opens the underlying bidi-gRPC object for reading. @@ -145,8 +146,8 @@ async def create_mrd( :type object_name: str :param object_name: The name of the object to be read. - :type generation_number: int - :param generation_number: (Optional) If present, selects a specific + :type generation: int + :param generation: (Optional) If present, selects a specific revision of this object. :type read_handle: _storage_v2.BidiReadHandle @@ -162,7 +163,14 @@ async def create_mrd( :rtype: :class:`~google.cloud.storage._experimental.asyncio.async_multi_range_downloader.AsyncMultiRangeDownloader` :returns: An initialized AsyncMultiRangeDownloader instance for reading. """ - mrd = cls(client, bucket_name, object_name, generation_number, read_handle) + mrd = cls( + client, + bucket_name, + object_name, + generation=generation, + read_handle=read_handle, + **kwargs, + ) await mrd.open(retry_policy=retry_policy, metadata=metadata) return mrd @@ -171,8 +179,9 @@ def __init__( client: AsyncGrpcClient, bucket_name: str, object_name: str, - generation_number: Optional[int] = None, + generation: Optional[int] = None, read_handle: Optional[_storage_v2.BidiReadHandle] = None, + **kwargs, ) -> None: """Constructor for AsyncMultiRangeDownloader, clients are not adviced to use it directly. Instead it's adviced to use the classmethod `create_mrd`. @@ -186,20 +195,27 @@ def __init__( :type object_name: str :param object_name: The name of the object to be read. - :type generation_number: int - :param generation_number: (Optional) If present, selects a specific revision of + :type generation: int + :param generation: (Optional) If present, selects a specific revision of this object. :type read_handle: _storage_v2.BidiReadHandle :param read_handle: (Optional) An existing read handle. """ + if "generation_number" in kwargs: + if generation is not None: + raise TypeError( + "Cannot set both 'generation' and 'generation_number'. " + "Use 'generation' for new code." + ) + generation = kwargs.pop("generation_number") raise_if_no_fast_crc32c() self.client = client self.bucket_name = bucket_name self.object_name = object_name - self.generation_number = generation_number + self.generation = generation self.read_handle: Optional[_storage_v2.BidiReadHandle] = read_handle self.read_obj_str: Optional[_AsyncReadObjectStream] = None self._is_stream_open: bool = False @@ -276,7 +292,7 @@ async def _do_open(): client=self.client.grpc_client, bucket_name=self.bucket_name, object_name=self.object_name, - generation_number=self.generation_number, + generation_number=self.generation, read_handle=self.read_handle, ) @@ -291,7 +307,7 @@ async def _do_open(): ) if self.read_obj_str.generation_number: - self.generation_number = self.read_obj_str.generation_number + self.generation = self.read_obj_str.generation_number if self.read_obj_str.read_handle: self.read_handle = self.read_obj_str.read_handle if self.read_obj_str.persisted_size is not None: @@ -435,7 +451,7 @@ async def generator(): client=self.client.grpc_client, bucket_name=self.bucket_name, object_name=self.object_name, - generation_number=self.generation_number, + generation_number=self.generation, read_handle=current_handle, ) diff --git a/tests/system/test_zonal.py b/tests/system/test_zonal.py index 4d46353e2..129462713 100644 --- a/tests/system/test_zonal.py +++ b/tests/system/test_zonal.py @@ -264,7 +264,8 @@ async def _run(): event_loop.run_until_complete(_run()) -@pytest.mark.skip(reason='Flaky test b/478129078') + +@pytest.mark.skip(reason="Flaky test b/478129078") def test_mrd_open_with_read_handle(event_loop, grpc_client_direct): object_name = f"test_read_handl-{str(uuid.uuid4())[:4]}" diff --git a/tests/unit/asyncio/test_async_multi_range_downloader.py b/tests/unit/asyncio/test_async_multi_range_downloader.py index 09cf0fa09..a1df97fda 100644 --- a/tests/unit/asyncio/test_async_multi_range_downloader.py +++ b/tests/unit/asyncio/test_async_multi_range_downloader.py @@ -49,7 +49,7 @@ async def _make_mock_mrd( mock_cls_async_read_object_stream, bucket_name=_TEST_BUCKET_NAME, object_name=_TEST_OBJECT_NAME, - generation_number=_TEST_GENERATION_NUMBER, + generation=_TEST_GENERATION_NUMBER, read_handle=_TEST_READ_HANDLE, ): mock_client = mock.MagicMock() @@ -62,7 +62,7 @@ async def _make_mock_mrd( mock_stream.read_handle = _TEST_READ_HANDLE mrd = await AsyncMultiRangeDownloader.create_mrd( - mock_client, bucket_name, object_name, generation_number, read_handle + mock_client, bucket_name, object_name, generation, read_handle ) return mrd, mock_client @@ -89,7 +89,7 @@ async def test_create_mrd(self, mock_cls_async_read_object_stream): assert mrd.client == mock_client assert mrd.bucket_name == _TEST_BUCKET_NAME assert mrd.object_name == _TEST_OBJECT_NAME - assert mrd.generation_number == _TEST_GENERATION_NUMBER + assert mrd.generation == _TEST_GENERATION_NUMBER assert mrd.read_handle == _TEST_READ_HANDLE assert mrd.persisted_size == _TEST_OBJECT_SIZE assert mrd.is_stream_open @@ -303,9 +303,7 @@ async def test_downloading_without_opening_should_throw_error(self): assert not mrd.is_stream_open @mock.patch("google.cloud.storage._experimental.asyncio._utils.google_crc32c") - def test_init_raises_if_crc32c_c_extension_is_missing( - self, mock_google_crc32c - ): + def test_init_raises_if_crc32c_c_extension_is_missing(self, mock_google_crc32c): mock_google_crc32c.implementation = "python" mock_client = mock.MagicMock() @@ -405,3 +403,47 @@ async def close_side_effect(): mock_close.assert_called_once() assert not mrd.is_stream_open + + @mock.patch( + "google.cloud.storage._experimental.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_create_mrd_with_generation_number( + self, mock_cls_async_read_object_stream + ): + # Arrange + mock_client = mock.MagicMock() + mock_client.grpc_client = mock.AsyncMock() + + mock_stream = mock_cls_async_read_object_stream.return_value + mock_stream.open = AsyncMock() + mock_stream.generation_number = _TEST_GENERATION_NUMBER + mock_stream.persisted_size = _TEST_OBJECT_SIZE + mock_stream.read_handle = _TEST_READ_HANDLE + + # Act + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, + _TEST_BUCKET_NAME, + _TEST_OBJECT_NAME, + generation_number=_TEST_GENERATION_NUMBER, + read_handle=_TEST_READ_HANDLE, + ) + + # Assert + assert mrd.generation == _TEST_GENERATION_NUMBER + + @pytest.mark.asyncio + async def test_create_mrd_with_both_generation_and_generation_number(self): + # Arrange + mock_client = mock.MagicMock() + + # Act & Assert + with pytest.raises(TypeError): + await AsyncMultiRangeDownloader.create_mrd( + mock_client, + _TEST_BUCKET_NAME, + _TEST_OBJECT_NAME, + generation=_TEST_GENERATION_NUMBER, + generation_number=_TEST_GENERATION_NUMBER, + )