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