# Copyright 2022 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. import dataclasses import enum import typing from unittest import mock import pytest from google.api_core import exceptions from google.api_core import extended_operation from google.api_core import retry TEST_OPERATION_NAME = "test/extended_operation" @dataclasses.dataclass(frozen=True) class CustomOperation: class StatusCode(enum.Enum): UNKNOWN = 0 DONE = 1 PENDING = 2 class LROCustomErrors: class LROCustomError: def __init__(self, code: str = "", message: str = ""): self.code = code self.message = message def __init__(self, errors: typing.List[LROCustomError] = []): self.errors = errors name: str status: StatusCode error_code: typing.Optional[int] = None error_message: typing.Optional[str] = None armor_class: typing.Optional[int] = None # Note: `error` can be removed once proposal A from # b/284179390 is implemented. error: typing.Optional[LROCustomErrors] = None # Note: in generated clients, this property must be generated for each # extended operation message type. # The status may be an enum, a string, or a bool. If it's a string or enum, # its text is compared to the string "DONE". @property def done(self): return self.status.name == "DONE" def make_extended_operation(responses=None): client_operations_responses = responses or [ CustomOperation( name=TEST_OPERATION_NAME, status=CustomOperation.StatusCode.PENDING ) ] refresh = mock.Mock(spec=["__call__"], side_effect=client_operations_responses) refresh.responses = client_operations_responses cancel = mock.Mock(spec=["__call__"]) extended_operation_future = extended_operation.ExtendedOperation.make( refresh, cancel, client_operations_responses[0], ) return extended_operation_future, refresh, cancel def test_constructor(): ex_op, refresh, _ = make_extended_operation() assert ex_op._extended_operation == refresh.responses[0] assert not ex_op.cancelled() assert not ex_op.done() assert ex_op.name == TEST_OPERATION_NAME assert ex_op.status == CustomOperation.StatusCode.PENDING assert ex_op.error_code is None assert ex_op.error_message is None def test_done(): responses = [ CustomOperation( name=TEST_OPERATION_NAME, status=CustomOperation.StatusCode.PENDING ), # Second response indicates that the operation has finished. CustomOperation( name=TEST_OPERATION_NAME, status=CustomOperation.StatusCode.DONE ), # Bumper to make sure we stop polling on DONE. CustomOperation( name=TEST_OPERATION_NAME, status=CustomOperation.StatusCode.DONE, error_message="Gone too far!", ), ] ex_op, refresh, _ = make_extended_operation(responses) # Start out not done. assert not ex_op.done() assert refresh.call_count == 1 # Refresh brings us to the done state. assert ex_op.done() assert refresh.call_count == 2 assert not ex_op.error_message # Make sure that subsequent checks are no-ops. assert ex_op.done() assert refresh.call_count == 2 assert not ex_op.error_message def test_cancellation(): responses = [ CustomOperation( name=TEST_OPERATION_NAME, status=CustomOperation.StatusCode.PENDING ), # Second response indicates that the operation was cancelled. CustomOperation( name=TEST_OPERATION_NAME, status=CustomOperation.StatusCode.DONE ), ] ex_op, _, cancel = make_extended_operation(responses) assert not ex_op.cancelled() assert ex_op.cancel() assert ex_op.cancelled() cancel.assert_called_once_with() # Cancelling twice should have no effect. assert not ex_op.cancel() cancel.assert_called_once_with() def test_done_w_retry(): # Not sure what's going on here with the coverage, so just ignore it. test_retry = retry.Retry(predicate=lambda x: True) # pragma: NO COVER responses = [ CustomOperation( name=TEST_OPERATION_NAME, status=CustomOperation.StatusCode.PENDING ), CustomOperation( name=TEST_OPERATION_NAME, status=CustomOperation.StatusCode.DONE ), ] ex_op, refresh, _ = make_extended_operation(responses) ex_op.done(retry=test_retry) refresh.assert_called_once_with(retry=test_retry) def test_error(): responses = [ CustomOperation( name=TEST_OPERATION_NAME, status=CustomOperation.StatusCode.DONE, error_code=400, error_message="Bad request", ), ] ex_op, _, _ = make_extended_operation(responses) # Defaults to CallError when grpc is not installed with pytest.raises(exceptions.BadRequest): ex_op.result() # Test GCE custom LRO Error. See b/284179390 # Note: This test case can be removed once proposal A from # b/284179390 is implemented. _EXCEPTION_CODE = "INCOMPATIBLE_BACKEND_SERVICES" _EXCEPTION_MESSAGE = "Validation failed for instance group" responses = [ CustomOperation( name=TEST_OPERATION_NAME, status=CustomOperation.StatusCode.DONE, error_code=400, error_message="Bad request", error=CustomOperation.LROCustomErrors( errors=[ CustomOperation.LROCustomErrors.LROCustomError( code=_EXCEPTION_CODE, message=_EXCEPTION_MESSAGE ) ] ), ), ] ex_op, _, _ = make_extended_operation(responses) # Defaults to CallError when grpc is not installed with pytest.raises( exceptions.BadRequest, match=f"{_EXCEPTION_CODE}: {_EXCEPTION_MESSAGE}" ): ex_op.result() # Inconsistent result responses = [ CustomOperation( name=TEST_OPERATION_NAME, status=CustomOperation.StatusCode.DONE, error_code=2112, ), ] ex_op, _, _ = make_extended_operation(responses) with pytest.raises(exceptions.GoogleAPICallError): ex_op.result() def test_pass_through(): responses = [ CustomOperation( name=TEST_OPERATION_NAME, status=CustomOperation.StatusCode.PENDING, armor_class=10, ), CustomOperation( name=TEST_OPERATION_NAME, status=CustomOperation.StatusCode.DONE, armor_class=20, ), ] ex_op, _, _ = make_extended_operation(responses) assert ex_op.armor_class == 10 ex_op.result() assert ex_op.armor_class == 20