def _set_result_from_operation(self):
"""Set the result or exception from the operation if it is complete."""
# This must be done in a lock to prevent the polling thread
# and main thread from both executing the completion logic
# at the same time.
with self._completion_lock:
# If the operation isn't complete or if the result has already been
# set, do not call set_result/set_exception again.
# Note: self._result_set is set to True in set_result and
# set_exception, in case those methods are invoked directly.
if not self._operation.done or self._result_set:
return
if self._operation.HasField("response"):
response = protobuf_helpers.from_any_pb(
self._result_type, self._operation.response)
self.set_result(response)
elif self._operation.HasField("error"):
exception = exceptions.GoogleAPICallError(
self._operation.error.message,
errors=(self._operation.error, ),
response=self._operation,
)
self.set_exception(exception)
else:
exception = exceptions.GoogleAPICallError(
"Unexpected state: Long-running operation had neither "
"response nor error set.")
self.set_exception(exception)
def test_create_google_cloud_error():
exception = exceptions.GoogleAPICallError('Testing')
exception.code = 600
assert str(exception) == '600 Testing'
assert exception.message == 'Testing'
assert exception.errors == []
assert exception.response is None
def _set_result_from_operation(self):
"""Set the result or exception from the operation if it is complete."""
# This must be done in a lock to prevent the async_future thread
# and main thread from both executing the completion logic
# at the same time.
with self._completion_lock:
# If the operation isn't complete or if the result has already been
# set, do not call set_result/set_exception again.
if not self._operation.done or self._future.done():
return
if self._operation.HasField("response"):
response = protobuf_helpers.from_any_pb(
self._result_type, self._operation.response
)
self.set_result(response)
elif self._operation.HasField("error"):
exception = exceptions.GoogleAPICallError(
self._operation.error.message,
errors=(self._operation.error,),
response=self._operation,
)
self.set_exception(exception)
else:
# Some APIs set `done: true`, with an empty response.
# Set the result to an empty message of the expected
# result type.
# https://google.aip.dev/151
self.set_result(self._result_type())
def test_send_unary_api_call_error(caplog):
caplog.set_level(logging.DEBUG)
manager = make_manager()
error = exceptions.GoogleAPICallError("The front fell off")
manager._client.acknowledge.side_effect = error
manager.send(gapic_types.StreamingPullRequest(ack_ids=["ack_id1", "ack_id2"]))
assert "The front fell off" in caplog.text
def test_send_unary_error(caplog):
caplog.set_level(logging.DEBUG)
manager = make_manager()
manager._UNARY_REQUESTS = True
error = exceptions.GoogleAPICallError('The front fell off')
manager._client.acknowledge.side_effect = error
manager.send(types.StreamingPullRequest(ack_ids=['ack_id1', 'ack_id2']))
assert 'The front fell off' in caplog.text
def test_create_google_cloud_error_with_args():
error = {
"domain": "global",
"location": "test",
"locationType": "testing",
"message": "Testing",
"reason": "test",
}
response = mock.sentinel.response
exception = exceptions.GoogleAPICallError("Testing", [error], response=response)
exception.code = 600
assert str(exception) == "600 Testing"
assert exception.message == "Testing"
assert exception.errors == [error]
assert exception.response == response
def test_create_google_cloud_error_with_args():
error = {
'domain': 'global',
'location': 'test',
'locationType': 'testing',
'message': 'Testing',
'reason': 'test',
}
response = mock.sentinel.response
exception = exceptions.GoogleAPICallError('Testing', [error],
response=response)
exception.code = 600
assert str(exception) == '600 Testing'
assert exception.message == 'Testing'
assert exception.errors == [error]
assert exception.response == response
class WaitForCompletionOperatorTest(unittest.TestCase):
def setUp(self):
super(WaitForCompletionOperatorTest, self).setUp()
dag = models.DAG(dag_id=_DAG_ID, start_date=datetime.datetime.now())
self._task = wait_for_completion_operator.WaitForCompletionOperator(
project_id=_PROJECT_ID,
queue_location=_QUEUE_LOCATION,
queue_name=_QUEUE_NAME,
service_account_path=_SERVICE_ACCOUNT_PATH,
try_count_limit=_TRY_COUNT_LIMIT,
dag=dag,
task_id=_TASK_ID)
self._context = mock.MagicMock()
self._mock_tasks_client = mock.patch(
'google.cloud.tasks.CloudTasksClient.from_service_account_json',
autospec=True).start()
self._mock_time_sleep = mock.patch('time.sleep').start()
self.addCleanup(mock.patch.stopall)
@parameterized.expand([([['dummy_task_0', 'dummy_task_1', 'dummy_task_2'],
['dummy_task_0'], []], 2), ([[]], 0)])
def test_execute_should_wait_until_queue_becomes_empty(
self, sequence_of_tasks, expected_sleep_count):
self._mock_tasks_client.return_value.list_tasks.side_effect = sequence_of_tasks
self._task.execute(self._context)
self.assertEqual(expected_sleep_count,
self._mock_time_sleep.call_count)
@parameterized.expand([[exceptions.GoogleAPICallError('')],
[exceptions.RetryError('', Exception())],
[ValueError()]])
def test_execute_when_exception_raised(self, api_exception):
self._mock_tasks_client.return_value.list_tasks.side_effect = api_exception
with self.assertRaises(airflow.AirflowException):
self._task.execute(self._context)
def test_execute_should_raise_exception_when_try_count_exceeds_limit(self):
self._mock_tasks_client.return_value.list_tasks.return_value = [
'blocked_task'
]
with self.assertRaises(airflow.AirflowException):
self._task.execute(self._context)
self.assertEqual(
_TRY_COUNT_LIMIT,
self._mock_tasks_client.return_value.list_tasks.call_count)
self.assertEqual(_TRY_COUNT_LIMIT, self._mock_time_sleep.call_count)
def test_get_entries_err_from_search_results_should_return_values(self):
entry = MockedObject()
entry.name = 'asset_1'
entry.relative_resource_name = 'asset_1_resource_name'
entry_2 = MockedObject()
entry_2.name = 'asset_2'
entry_2.relative_resource_name = 'asset_2_resource_name'
search_results = [entry, entry_2]
datacatalog = self.__datacatalog_client
datacatalog.get_entry.side_effect = [
{}, exceptions.GoogleAPICallError('Permission denied')
]
self.__datacatalog_facade.get_entries_from_search_results(
search_results)
self.assertEqual(2, datacatalog.get_entry.call_count)
def test_error_when_deleting_entry_group_should_be_ignored(self):
entries = []
datacatalog_facade = self.__datacatalog_facade
datacatalog_facade.search_catalog_relative_resource_name\
.return_value = [
'projects/uat-env-1/locations/us-central1/'
'entryGroups/system/entries/database',
'projects/uat-env-1/locations/us-central1/'
'entryGroups/system/entries/table'
]
datacatalog_facade.delete_entry_group.side_effect = \
exceptions.GoogleAPICallError('Error when deleting entry group')
self.__metadata_cleaner.delete_obsolete_metadata(
entries, self.__SEARCH_QUERY)
self.assertEqual(
1, datacatalog_facade.search_catalog_relative_resource_name.
call_count)
self.assertEqual(2, datacatalog_facade.delete_entry.call_count)
self.assertEqual(1, datacatalog_facade.delete_entry_group.call_count)
def test_w_google_api_call_error_miss(self):
from google.api_core import exceptions
exc = exceptions.GoogleAPICallError("testing")
exc.code = 999
self.assertFalse(self._call_fut(exc))
def test_w_google_api_call_error_hit(self):
from google.api_core import exceptions
exc = exceptions.GoogleAPICallError("testing")
exc.code = 408
self.assertTrue(self._call_fut(exc))
