# 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

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