def test_operations_contain_expected_statuses_red_path(self, statuses, expected_statuses):
operations = [{NAME: TEST_TRANSFER_OPERATION_NAME, METADATA: {STATUS: status}} for status in statuses]
with self.assertRaisesRegex(
AirflowException,
"An unexpected operation status was encountered. Expected: {}".format(
", ".join(expected_statuses)
),
):
GCPTransferServiceHook.operations_contain_expected_statuses(
operations, GcpTransferOperationStatus.IN_PROGRESS
)
def test_operations_contain_expected_statuses_green_path(
self, statuses, expected_statuses):
operations = [{
NAME: TEST_TRANSFER_OPERATION_NAME,
METADATA: {
STATUS: status
}
} for status in statuses]
result = GCPTransferServiceHook.operations_contain_expected_statuses(
operations, expected_statuses)
self.assertTrue(result)
def poke(self, context):
hook = GCPTransferServiceHook(gcp_conn_id=self.gcp_cloud_conn_id)
operations = hook.list_transfer_operations(
request_filter={'project_id': self.project_id, 'job_names': [self.job_name]}
)
check = GCPTransferServiceHook.operations_contain_expected_statuses(
operations=operations, expected_statuses=self.expected_statuses
)
if check:
self.xcom_push(key="sensed_operations", value=operations, context=context)
return check