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) ), ): CloudDataTransferServiceHook.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 = CloudDataTransferServiceHook.operations_contain_expected_statuses( operations, expected_statuses ) self.assertTrue(result)
def poke(self, context): hook = CloudDataTransferServiceHook(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 = CloudDataTransferServiceHook.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
def poke(self, context: dict) -> bool: hook = CloudDataTransferServiceHook( gcp_conn_id=self.gcp_cloud_conn_id, impersonation_chain=self.impersonation_chain, ) operations = hook.list_transfer_operations(request_filter={ 'project_id': self.project_id, 'job_names': [self.job_name] }) for operation in operations: self.log.info("Progress for operation %s: %s", operation[NAME], operation[METADATA][COUNTERS]) check = CloudDataTransferServiceHook.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