diff --git a/google/cloud/bigtable/admin_v2/overlay/services/bigtable_table_admin/async_client.py b/google/cloud/bigtable/admin_v2/overlay/services/bigtable_table_admin/async_client.py index bc96ba9b2..809b4614b 100644 --- a/google/cloud/bigtable/admin_v2/overlay/services/bigtable_table_admin/async_client.py +++ b/google/cloud/bigtable/admin_v2/overlay/services/bigtable_table_admin/async_client.py @@ -52,6 +52,7 @@ ) from google.cloud.bigtable.admin_v2.overlay.types import ( async_consistency, + async_restore_table, wait_for_consistency_request, ) @@ -133,6 +134,100 @@ def __init__( client_info=client_info, ) + async def restore_table( + self, + request: Optional[Union[bigtable_table_admin.RestoreTableRequest, dict]] = None, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> async_restore_table.AsyncRestoreTableOperation: + r"""Create a new table by restoring from a completed backup. The + returned table :class:`long-running operation + ` + can be used to track the progress of the operation, and to cancel it. The + :attr:`metadata ` field type is + :class:`RestoreTableMetadata `. + The :meth:`response ` type is + :class:`google.cloud.bigtable.admin_v2.types.Table`, if successful. + + Additionally, the returned :class:`long-running-operation ` + provides a method, :meth:`google.cloud.bigtable.admin_v2.overlay.types.async_restore_table.AsyncRestoreTableOperation.optimize_restore_table_operation` that + provides access to a :class:`google.api_core.operation_async.AsyncOperation` object representing the OptimizeRestoreTable long-running-operation + after the current one has completed. + + .. code-block:: python + + # This snippet should be regarded as a code template only. + # + # It will require modifications to work: + # - It may require correct/in-range values for request initialization. + # - It may require specifying regional endpoints when creating the service + # client as shown in: + # https://googleapis.dev/python/google-api-core/latest/client_options.html + from google.cloud.bigtable import admin_v2 + + async def sample_restore_table(): + # Create a client + client = admin_v2.BigtableTableAdminAsyncClient() + + # Initialize request argument(s) + request = admin_v2.RestoreTableRequest( + backup="backup_value", + parent="parent_value", + table_id="table_id_value", + ) + + # Make the request + operation = await client.restore_table(request=request) + + print("Waiting for operation to complete...") + + response = await operation.result() + + # Handle the response + print(response) + + # Handle LRO2 + optimize_operation = await operation.optimize_restore_table_operation() + + if optimize_operation: + print("Waiting for table optimization to complete...") + + response = await optimize_operation.result() + + Args: + request (Union[google.cloud.bigtable.admin_v2.types.RestoreTableRequest, dict]): + The request object. The request for + [RestoreTable][google.bigtable.admin.v2.BigtableTableAdmin.RestoreTable]. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + google.cloud.bigtable.admin_v2.overlay.types.async_restore_table.AsyncRestoreTableOperation: + An object representing a long-running operation. + + The result type for the operation will be :class:`google.cloud.bigtable.admin_v2.types.Table` A collection of user data indexed by row, column, and timestamp. + Each table is served using the resources of its + parent cluster. + """ + operation = await self._restore_table( + request=request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + restore_table_operation = async_restore_table.AsyncRestoreTableOperation( + self._client._transport.operations_client, operation + ) + return restore_table_operation + async def wait_for_consistency( self, request: Optional[ diff --git a/google/cloud/bigtable/admin_v2/overlay/types/async_restore_table.py b/google/cloud/bigtable/admin_v2/overlay/types/async_restore_table.py new file mode 100644 index 000000000..e371b6425 --- /dev/null +++ b/google/cloud/bigtable/admin_v2/overlay/types/async_restore_table.py @@ -0,0 +1,99 @@ +# Copyright 2025 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional + +from google.api_core import exceptions +from google.api_core import operation_async +from google.protobuf import empty_pb2 + +from google.cloud.bigtable.admin_v2.types import OptimizeRestoredTableMetadata + + +class AsyncRestoreTableOperation(operation_async.AsyncOperation): + """A Future for interacting with Bigtable Admin's RestoreTable Long-Running Operation. + + This is needed to expose a potential long-running operation that might run after this operation + finishes, OptimizeRestoreTable. This is exposed via the the :meth:`optimize_restore_table_operation` + method. + + **This class should not be instantiated by users** and should only be instantiated by the admin + client's :meth:`restore_table + ` + method. + + Args: + operations_client (google.api_core.operations_v1.AbstractOperationsClient): The operations + client from the admin client class's transport. + restore_table_operation (google.api_core.operation_async.AsyncOperation): A + :class:`google.api_core.operation_async.AsyncOperation` + instance resembling a RestoreTable long-running operation + """ + + def __init__( + self, operations_client, restore_table_operation: operation_async.AsyncOperation + ): + self._operations_client = operations_client + self._optimize_restored_table_operation = None + super().__init__( + restore_table_operation._operation, + restore_table_operation._refresh, + restore_table_operation._cancel, + restore_table_operation._result_type, + restore_table_operation._metadata_type, + retry=restore_table_operation._retry, + ) + + async def optimize_restored_table_operation( + self, + ) -> Optional[operation_async.AsyncOperation]: + """Gets the OptimizeRestoredTable long-running operation that runs after this operation finishes. + The current operation might not trigger a follow-up OptimizeRestoredTable operation, in which case, this + method will return `None`. + This method must not be called before the parent restore_table operation is complete. + Returns: + An object representing a long-running operation, or None if there is no OptimizeRestoredTable operation + after this one. + Raises: + RuntimeError: raised when accessed before the restore_table operation is complete + + Raises: + google.api_core.GoogleAPIError: raised when accessed before the restore_table operation is complete + """ + if not await self.done(): + raise exceptions.GoogleAPIError( + "optimize_restored_table operation can't be accessed until the restore_table operation is complete" + ) + + if self._optimize_restored_table_operation is not None: + return self._optimize_restored_table_operation + + operation_name = self.metadata.optimize_table_operation_name + + # When the RestoreTable operation finishes, it might not necessarily trigger + # an optimize operation. + if operation_name: + gapic_operation = await self._operations_client.get_operation( + name=operation_name + ) + self._optimize_restored_table_operation = operation_async.from_gapic( + gapic_operation, + self._operations_client, + empty_pb2.Empty, + metadata_type=OptimizeRestoredTableMetadata, + ) + return self._optimize_restored_table_operation + else: + # no optimize operation found + return None diff --git a/google/cloud/bigtable/admin_v2/overlay/types/restore_table.py b/google/cloud/bigtable/admin_v2/overlay/types/restore_table.py index c346a6b79..aae3a4bf4 100644 --- a/google/cloud/bigtable/admin_v2/overlay/types/restore_table.py +++ b/google/cloud/bigtable/admin_v2/overlay/types/restore_table.py @@ -12,22 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -import enum -from typing import Optional, Union +from typing import Optional -from google.api_core import retry +from google.api_core import exceptions from google.api_core import operation from google.protobuf import empty_pb2 from google.cloud.bigtable.admin_v2.types import OptimizeRestoredTableMetadata -# This is needed due to a documentation issue with using object() as the default -# value for a parameter. -class Timeout(enum.Enum): - DEFAULT_TIMEOUT = "DEFAULT_TIMEOUT" - - class RestoreTableOperation(operation.Operation): """A Future for interacting with Bigtable Admin's RestoreTable Long-Running Operation. @@ -59,16 +52,14 @@ def __init__(self, operations_client, restore_table_operation: operation.Operati polling=restore_table_operation._polling, ) - def optimize_restored_table_operation( - self, - timeout: Optional[Union[int, Timeout]] = Timeout.DEFAULT_TIMEOUT, - retry: Optional[retry.Retry] = None, - polling: Optional[retry.Retry] = None, - ) -> Optional[operation.Operation]: + def optimize_restored_table_operation(self) -> Optional[operation.Operation]: """Gets the OptimizeRestoredTable long-running operation that runs after this operation finishes. - This is a blocking call that will return the operation after this current long-running operation - finishes, just like :meth:`google.api_core.operation.Operation.result`. The follow-up operation has + This must not be called before the parent restore_table operation is complete. You can guarantee + this happening by calling this function after this class's :meth:`google.api_core.operation.Operation.result` + method. + + The follow-up operation has :attr:`metadata ` type :class:`OptimizeRestoredTableMetadata ` @@ -77,51 +68,35 @@ def optimize_restored_table_operation( The current operation might not trigger a follow-up OptimizeRestoredTable operation, in which case, this method will return `None`. - Args: - timeout (Optional[int | google.cloud.bigtable.admin_v2.overlay.types.restore_table.Timeout]): - How long (in seconds) to wait for the operation to complete. If None, wait indefinitely. If - `Timeout.DEFAULT_TIMEOUT`, wait the default amount. - retry (Optional[google.api_core.retry.Retry]): How to retry the polling RPC. This defines ONLY - how the polling RPC call is retried (i.e. what to do if the RPC we used for polling returned - an error). It does NOT define how the polling is done (i.e. how frequently and for how long - to call the polling RPC). - polling (Optional[google.api_core.retry.Retry]): How often and for how long to call polling RPC - periodically. This parameter does NOT define how to retry each individual polling RPC call - (use the `retry` parameter for that). - Returns: Optional[google.api_core.operation.Operation]: An object representing a long-running operation, or None if there is no OptimizeRestoredTable operation after this one. - """ - if timeout == Timeout.DEFAULT_TIMEOUT: - timeout = operation.Operation._DEFAULT_VALUE - - self._blocking_poll(timeout=timeout, retry=retry, polling=polling) - if self._exception is not None: - # pylint: disable=raising-bad-type - # Pylint doesn't recognize that this is valid in this case. - raise self._exception + Raises: + google.api_core.GoogleAPIError: raised when accessed before the restore_table operation is complete + """ + if not self.done(): + raise exceptions.GoogleAPIError( + "optimize_restored_table operation can't be accessed until the restore_table operation is complete" + ) - return self._optimize_restored_table_operation + if self._optimize_restored_table_operation is not None: + return self._optimize_restored_table_operation - def set_result(self, response): - optimize_restored_table_operation_name = ( - self.metadata.optimize_table_operation_name - ) + operation_name = self.metadata.optimize_table_operation_name # When the RestoreTable operation finishes, it might not necessarily trigger # an optimize operation. - if optimize_restored_table_operation_name: - optimize_restore_table_operation = self._operations_client.get_operation( - name=optimize_restored_table_operation_name - ) + if operation_name: + gapic_operation = self._operations_client.get_operation(name=operation_name) self._optimize_restored_table_operation = operation.from_gapic( - optimize_restore_table_operation, + gapic_operation, self._operations_client, empty_pb2.Empty, metadata_type=OptimizeRestoredTableMetadata, ) - - super().set_result(response) + return self._optimize_restored_table_operation + else: + # no optimize operation found + return None diff --git a/tests/unit/admin_overlay/test_async_client.py b/tests/unit/admin_overlay/test_async_client.py index cf83e3813..c31e68813 100644 --- a/tests/unit/admin_overlay/test_async_client.py +++ b/tests/unit/admin_overlay/test_async_client.py @@ -16,7 +16,7 @@ # try/except added for compatibility with python < 3.8 try: from unittest import mock - from unittest.mock import AsyncMock # pragma: NO COVER + from unittest.mock import AsyncMock # pragma: NO COVER # noqa: F401 except ImportError: # pragma: NO COVER import mock @@ -29,7 +29,10 @@ BigtableTableAdminAsyncClient, DEFAULT_CLIENT_INFO, ) -from google.cloud.bigtable.admin_v2.overlay.types import wait_for_consistency_request +from google.cloud.bigtable.admin_v2.overlay.types import ( + async_restore_table, + wait_for_consistency_request, +) from google.cloud.bigtable import __version__ as bigtable_version @@ -73,6 +76,60 @@ def test_bigtable_table_admin_async_client_client_version( ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "kwargs", + [ + { + "request": bigtable_table_admin.RestoreTableRequest( + parent=PARENT_NAME, + table_id=TABLE_NAME, + ) + }, + { + "request": { + "parent": PARENT_NAME, + "table_id": TABLE_NAME, + }, + }, + { + "request": bigtable_table_admin.RestoreTableRequest( + parent=PARENT_NAME, + table_id=TABLE_NAME, + ), + "retry": mock.Mock(spec=retries.Retry), + "timeout": mock.Mock(spec=retries.Retry), + "metadata": [("foo", "bar")], + }, + ], +) +async def test_bigtable_table_admin_async_client_restore_table(kwargs): + client = BigtableTableAdminAsyncClient() + + with mock.patch.object( + async_restore_table, "AsyncRestoreTableOperation", new_callable=mock.AsyncMock + ) as future_mock: + with mock.patch.object( + client._client, "_transport", new_callable=mock.AsyncMock + ) as transport_mock: + with mock.patch.object( + client, "_restore_table", new_callable=mock.AsyncMock + ) as restore_table_mock: + operation_mock = mock.Mock() + restore_table_mock.return_value = operation_mock + await client.restore_table(**kwargs) + + restore_table_mock.assert_called_once_with( + request=kwargs["request"], + retry=kwargs.get("retry", gapic_v1.method.DEFAULT), + timeout=kwargs.get("timeout", gapic_v1.method.DEFAULT), + metadata=kwargs.get("metadata", ()), + ) + future_mock.assert_called_once_with( + transport_mock.operations_client, operation_mock + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( "kwargs,check_consistency_request_extras", diff --git a/tests/unit/admin_overlay/test_async_restore_table.py b/tests/unit/admin_overlay/test_async_restore_table.py new file mode 100644 index 000000000..e3f01ee5a --- /dev/null +++ b/tests/unit/admin_overlay/test_async_restore_table.py @@ -0,0 +1,248 @@ +# Copyright 2025 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# try/except added for compatibility with python < 3.8 +try: + from unittest import mock + from unittest.mock import AsyncMock # pragma: NO COVER # noqa: F401 +except ImportError: # pragma: NO COVER + import mock + +from google.longrunning import operations_pb2 +from google.rpc import status_pb2, code_pb2 + +from google.api_core import operation_async, exceptions +from google.api_core.future import async_future +from google.api_core.operations_v1 import operations_async_client +from google.cloud.bigtable.admin_v2.types import bigtable_table_admin, table +from google.cloud.bigtable.admin_v2.overlay.types import async_restore_table + +import pytest + + +# Set up the mock operations +DEFAULT_MAX_POLL = 3 +RESTORE_TABLE_OPERATION_TABLE_NAME = "Test Table" +RESTORE_TABLE_OPERATION_NAME = "test/restore_table" +RESTORE_TABLE_OPERATION_METADATA = bigtable_table_admin.RestoreTableMetadata( + name=RESTORE_TABLE_OPERATION_TABLE_NAME, +) +OPTIMIZE_RESTORED_TABLE_OPERATION_NAME = "test/optimize_restore_table" +OPTIMIZE_RESTORED_TABLE_METADATA = bigtable_table_admin.OptimizeRestoredTableMetadata( + name=RESTORE_TABLE_OPERATION_TABLE_NAME, +) + +OPTIMIZE_RESTORED_TABLE_OPERATION_ID = "abcdefg" +RESTORE_TABLE_OPERATION_FINISHED_RESPONSE = table.Table( + name=RESTORE_TABLE_OPERATION_TABLE_NAME, +) +RESTORE_TABLE_OPERATION_FINISHED_ERROR = status_pb2.Status( + code=code_pb2.DEADLINE_EXCEEDED, message="Deadline Exceeded" +) + + +def make_operation_proto( + name, done=False, metadata=None, response=None, error=None, **kwargs +): + operation_proto = operations_pb2.Operation(name=name, done=done, **kwargs) + + if metadata is not None: + operation_proto.metadata.Pack(metadata._pb) + + if response is not None: + operation_proto.response.Pack(response._pb) + + if error is not None: + operation_proto.error.CopyFrom(error) + + return operation_proto + + +RESTORE_TABLE_IN_PROGRESS_OPERATION_PROTO = make_operation_proto( + name=RESTORE_TABLE_OPERATION_NAME, + done=False, + metadata=RESTORE_TABLE_OPERATION_METADATA, +) + +OPTIMIZE_RESTORED_TABLE_OPERATION_PROTO = make_operation_proto( + name=OPTIMIZE_RESTORED_TABLE_OPERATION_NAME, + metadata=OPTIMIZE_RESTORED_TABLE_METADATA, +) + + +# Set up the mock operation client +def mock_restore_table_operation( + max_poll_count=DEFAULT_MAX_POLL, fail=False, has_optimize_operation=True +): + client = mock.AsyncMock(spec=operations_async_client.OperationsAsyncClient) + + # Set up the polling + side_effect = [RESTORE_TABLE_IN_PROGRESS_OPERATION_PROTO] * (max_poll_count - 1) + finished_operation_metadata = bigtable_table_admin.RestoreTableMetadata() + bigtable_table_admin.RestoreTableMetadata.copy_from( + finished_operation_metadata, RESTORE_TABLE_OPERATION_METADATA + ) + if has_optimize_operation: + finished_operation_metadata.optimize_table_operation_name = ( + OPTIMIZE_RESTORED_TABLE_OPERATION_ID + ) + + if fail: + final_operation_proto = make_operation_proto( + name=RESTORE_TABLE_OPERATION_NAME, + done=True, + metadata=finished_operation_metadata, + error=RESTORE_TABLE_OPERATION_FINISHED_ERROR, + ) + else: + final_operation_proto = make_operation_proto( + name=RESTORE_TABLE_OPERATION_NAME, + done=True, + metadata=finished_operation_metadata, + response=RESTORE_TABLE_OPERATION_FINISHED_RESPONSE, + ) + side_effect.append(final_operation_proto) + refresh = mock.AsyncMock(spec=["__call__"], side_effect=side_effect) + cancel = mock.AsyncMock(spec=["__call__"]) + future = operation_async.AsyncOperation( + RESTORE_TABLE_IN_PROGRESS_OPERATION_PROTO, + refresh, + cancel, + result_type=table.Table, + metadata_type=bigtable_table_admin.RestoreTableMetadata, + ) + + # Set up the optimize_restore_table_operation + client.get_operation.side_effect = [OPTIMIZE_RESTORED_TABLE_OPERATION_PROTO] + + return async_restore_table.AsyncRestoreTableOperation(client, future) + + +@pytest.mark.asyncio +async def test_async_restore_table_operation_client_success_has_optimize(): + restore_table_operation = mock_restore_table_operation() + + await restore_table_operation.result() + optimize_restored_table_operation = ( + await restore_table_operation.optimize_restored_table_operation() + ) + + assert isinstance(optimize_restored_table_operation, operation_async.AsyncOperation) + assert ( + optimize_restored_table_operation._operation + == OPTIMIZE_RESTORED_TABLE_OPERATION_PROTO + ) + restore_table_operation._operations_client.get_operation.assert_called_with( + name=OPTIMIZE_RESTORED_TABLE_OPERATION_ID + ) + restore_table_operation._refresh.assert_has_calls( + [mock.call(retry=async_future.DEFAULT_RETRY)] * DEFAULT_MAX_POLL + ) + + +@pytest.mark.asyncio +async def test_restore_table_operation_client_success_has_optimize_multiple_calls(): + restore_table_operation = mock_restore_table_operation() + + await restore_table_operation.result() + optimize_restored_table_operation = ( + await restore_table_operation.optimize_restored_table_operation() + ) + + assert isinstance(optimize_restored_table_operation, operation_async.AsyncOperation) + assert ( + optimize_restored_table_operation._operation + == OPTIMIZE_RESTORED_TABLE_OPERATION_PROTO + ) + restore_table_operation._operations_client.get_operation.assert_called_with( + name=OPTIMIZE_RESTORED_TABLE_OPERATION_ID + ) + restore_table_operation._refresh.assert_has_calls( + [mock.call(retry=async_future.DEFAULT_RETRY)] * DEFAULT_MAX_POLL + ) + + await restore_table_operation.optimize_restored_table_operation() + restore_table_operation._refresh.assert_has_calls( + [mock.call(retry=async_future.DEFAULT_RETRY)] * DEFAULT_MAX_POLL + ) + + +@pytest.mark.asyncio +async def test_restore_table_operation_success_has_optimize_call_before_done(): + restore_table_operation = mock_restore_table_operation() + + with pytest.raises(exceptions.GoogleAPIError): + await restore_table_operation.optimize_restored_table_operation() + + restore_table_operation._operations_client.get_operation.assert_not_called() + + +@pytest.mark.asyncio +async def test_restore_table_operation_client_success_only_cache_after_finishing(): + restore_table_operation = mock_restore_table_operation() + + with pytest.raises(exceptions.GoogleAPIError): + await restore_table_operation.optimize_restored_table_operation() + + await restore_table_operation.result() + optimize_restored_table_operation = ( + await restore_table_operation.optimize_restored_table_operation() + ) + + assert isinstance(optimize_restored_table_operation, operation_async.AsyncOperation) + assert ( + optimize_restored_table_operation._operation + == OPTIMIZE_RESTORED_TABLE_OPERATION_PROTO + ) + restore_table_operation._operations_client.get_operation.assert_called_with( + name=OPTIMIZE_RESTORED_TABLE_OPERATION_ID + ) + restore_table_operation._refresh.assert_has_calls( + [mock.call(retry=async_future.DEFAULT_RETRY)] * DEFAULT_MAX_POLL + ) + + restore_table_operation.optimize_restored_table_operation() + restore_table_operation._refresh.assert_has_calls( + [mock.call(retry=async_future.DEFAULT_RETRY)] * DEFAULT_MAX_POLL + ) + + +@pytest.mark.asyncio +async def test_restore_table_operation_success_no_optimize(): + restore_table_operation = mock_restore_table_operation(has_optimize_operation=False) + + await restore_table_operation.result() + optimize_restored_table_operation = ( + await restore_table_operation.optimize_restored_table_operation() + ) + + assert optimize_restored_table_operation is None + restore_table_operation._operations_client.get_operation.assert_not_called() + + +@pytest.mark.asyncio +async def test_restore_table_operation_exception(): + restore_table_operation = mock_restore_table_operation( + fail=True, has_optimize_operation=False + ) + + with pytest.raises(exceptions.GoogleAPICallError): + await restore_table_operation.result() + + optimize_restored_table_operation = ( + await restore_table_operation.optimize_restored_table_operation() + ) + + assert optimize_restored_table_operation is None + restore_table_operation._operations_client.get_operation.assert_not_called() diff --git a/tests/unit/admin_overlay/test_restore_table.py b/tests/unit/admin_overlay/test_restore_table.py index f9cd2f4f0..75ba3a207 100644 --- a/tests/unit/admin_overlay/test_restore_table.py +++ b/tests/unit/admin_overlay/test_restore_table.py @@ -130,6 +130,7 @@ def mock_restore_table_operation( def test_restore_table_operation_client_success_has_optimize(): restore_table_operation = mock_restore_table_operation() + restore_table_operation.result() optimize_restored_table_operation = ( restore_table_operation.optimize_restored_table_operation() ) @@ -145,25 +146,43 @@ def test_restore_table_operation_client_success_has_optimize(): restore_table_operation._refresh.assert_has_calls([mock.call()] * DEFAULT_MAX_POLL) -@pytest.mark.parametrize( - "input_timeout,expected_timeout", - [ - (restore_table.Timeout.DEFAULT_TIMEOUT, operation.Operation._DEFAULT_VALUE), - (100, 100), - ], -) -def test_restore_table_timeouts(input_timeout, expected_timeout): +def test_restore_table_operation_client_success_has_optimize_multiple_calls(): restore_table_operation = mock_restore_table_operation() - with mock.patch.object(restore_table_operation, "_blocking_poll") as poll_mock: - restore_table_operation.optimize_restored_table_operation(timeout=input_timeout) - poll_mock.assert_called_once_with( - timeout=expected_timeout, retry=None, polling=None - ) + restore_table_operation.result() + optimize_restored_table_operation = ( + restore_table_operation.optimize_restored_table_operation() + ) + + assert isinstance(optimize_restored_table_operation, operation.Operation) + assert ( + optimize_restored_table_operation._operation + == OPTIMIZE_RESTORED_TABLE_OPERATION_PROTO + ) + restore_table_operation._operations_client.get_operation.assert_called_with( + name=OPTIMIZE_RESTORED_TABLE_OPERATION_ID + ) + restore_table_operation._refresh.assert_has_calls([mock.call()] * DEFAULT_MAX_POLL) + + restore_table_operation.optimize_restored_table_operation() + restore_table_operation._refresh.assert_has_calls([mock.call()] * DEFAULT_MAX_POLL) -def test_restore_table_operation_success_has_optimize_also_call_result(): + +def test_restore_table_operation_success_has_optimize_call_before_done(): restore_table_operation = mock_restore_table_operation() + with pytest.raises(exceptions.GoogleAPIError): + restore_table_operation.optimize_restored_table_operation() + + restore_table_operation._operations_client.get_operation.assert_not_called() + + +def test_restore_table_operation_client_success_only_cache_after_finishing(): + restore_table_operation = mock_restore_table_operation() + + with pytest.raises(exceptions.GoogleAPIError): + restore_table_operation.optimize_restored_table_operation() + restore_table_operation.result() optimize_restored_table_operation = ( restore_table_operation.optimize_restored_table_operation() @@ -179,6 +198,9 @@ def test_restore_table_operation_success_has_optimize_also_call_result(): ) restore_table_operation._refresh.assert_has_calls([mock.call()] * DEFAULT_MAX_POLL) + restore_table_operation.optimize_restored_table_operation() + restore_table_operation._refresh.assert_has_calls([mock.call()] * DEFAULT_MAX_POLL) + def test_restore_table_operation_success_no_optimize(): restore_table_operation = mock_restore_table_operation(has_optimize_operation=False) @@ -198,4 +220,11 @@ def test_restore_table_operation_exception(): ) with pytest.raises(exceptions.GoogleAPICallError): + restore_table_operation.result() + + optimize_restored_table_operation = ( restore_table_operation.optimize_restored_table_operation() + ) + + assert optimize_restored_table_operation is None + restore_table_operation._operations_client.get_operation.assert_not_called()