def execute(self, context): transfer_hook = GCPTransferServiceHook( gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) s3_creds = S3Hook(aws_conn_id=self.aws_conn_id).get_credentials() transfer_hook.create_transfer_job( project_id=self.project_id, transfer_spec={ 'awsS3DataSource': { 'bucketName': self.s3_bucket, 'awsAccessKey': { 'accessKeyId': s3_creds.access_key, 'secretAccessKey': s3_creds.secret_key, } }, 'gcsDataSink': { 'bucketName': self.gcs_bucket, }, 'objectConditions': self.object_conditions, 'transferOptions': self.transfer_options, }, **self.job_kwargs )
def execute(self, context): transfer_hook = GCPTransferServiceHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) s3_creds = S3Hook(aws_conn_id=self.aws_conn_id).get_credentials() transfer_hook.create_transfer_job(project_id=self.project_id, transfer_spec={ 'awsS3DataSource': { 'bucketName': self.s3_bucket, 'awsAccessKey': { 'accessKeyId': s3_creds.access_key, 'secretAccessKey': s3_creds.secret_key, } }, 'gcsDataSink': { 'bucketName': self.gcs_bucket, }, 'objectConditions': self.object_conditions, 'transferOptions': self.transfer_options, }, **self.job_kwargs)
def execute(self, context): hook = GCPTransferServiceHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) body = self._create_body() TransferJobPreprocessor(body=body, aws_conn_id=self.aws_conn_id).process_body() job = hook.create_transfer_job(body=body) if self.wait: hook.wait_for_transfer_job(job, timeout=self.timeout)
def execute(self, context): transfer_hook = GCPTransferServiceHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) job = transfer_hook.create_transfer_job( project_id=self.project_id, description=self.description, schedule=self.schedule, transfer_spec={ 'gcsDataSource': { 'bucketName': self.source_bucket, }, 'gcsDataSink': { 'bucketName': self.destination_bucket, }, 'objectConditions': self.object_conditions, 'transferOptions': self.transfer_options, }) if self.wait: transfer_hook.wait_for_transfer_job(job)
def execute(self, context): transfer_hook = GCPTransferServiceHook( gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) job = transfer_hook.create_transfer_job( project_id=self.project_id, description=self.description, schedule=self.schedule, transfer_spec={ 'gcsDataSource': { 'bucketName': self.source_bucket, }, 'gcsDataSink': { 'bucketName': self.destination_bucket, }, 'objectConditions': self.object_conditions, 'transferOptions': self.transfer_options, } ) if self.wait: transfer_hook.wait_for_transfer_job(job)
def execute(self, context): TransferJobPreprocessor(body=self.body, aws_conn_id=self.aws_conn_id).process_body() hook = GCPTransferServiceHook(api_version=self.api_version, gcp_conn_id=self.gcp_conn_id) return hook.create_transfer_job(body=self.body)
class TestGCPTransferServiceHookWithPassedProjectId(unittest.TestCase): def setUp(self): with mock.patch( 'airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__', new=mock_base_gcp_hook_no_default_project_id, ): self.gct_hook = GCPTransferServiceHook(gcp_conn_id='test') @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_create_transfer_job(self, get_conn): create_method = get_conn.return_value.transferJobs.return_value.create execute_method = create_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB res = self.gct_hook.create_transfer_job(body=TEST_BODY) self.assertEqual(res, TEST_TRANSFER_JOB) create_method.assert_called_once_with(body=TEST_BODY) execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_get_transfer_job(self, get_conn): get_method = get_conn.return_value.transferJobs.return_value.get execute_method = get_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB res = self.gct_hook.get_transfer_job(job_name=TEST_TRANSFER_JOB_NAME, project_id=TEST_PROJECT_ID) self.assertIsNotNone(res) self.assertEqual(TEST_TRANSFER_JOB_NAME, res[NAME]) get_method.assert_called_once_with(jobName=TEST_TRANSFER_JOB_NAME, projectId=TEST_PROJECT_ID) execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_list_transfer_job(self, get_conn): list_method = get_conn.return_value.transferJobs.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = {TRANSFER_JOBS: [TEST_TRANSFER_JOB]} list_next = get_conn.return_value.transferJobs.return_value.list_next list_next.return_value = None res = self.gct_hook.list_transfer_job( request_filter=TEST_TRANSFER_JOB_FILTER) self.assertIsNotNone(res) self.assertEqual(res, [TEST_TRANSFER_JOB]) list_method.assert_called_once_with(filter=mock.ANY) args, kwargs = list_method.call_args_list[0] self.assertEqual( json.loads(kwargs['filter']), { FILTER_PROJECT_ID: TEST_PROJECT_ID, FILTER_JOB_NAMES: [TEST_TRANSFER_JOB_NAME] }, ) list_execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_update_transfer_job(self, get_conn): update_method = get_conn.return_value.transferJobs.return_value.patch execute_method = update_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB res = self.gct_hook.update_transfer_job( job_name=TEST_TRANSFER_JOB_NAME, body=TEST_UPDATE_TRANSFER_JOB_BODY) self.assertIsNotNone(res) update_method.assert_called_once_with( jobName=TEST_TRANSFER_JOB_NAME, body=TEST_UPDATE_TRANSFER_JOB_BODY) execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_delete_transfer_job(self, get_conn): update_method = get_conn.return_value.transferJobs.return_value.patch execute_method = update_method.return_value.execute self.gct_hook.delete_transfer_job(job_name=TEST_TRANSFER_JOB_NAME, project_id=TEST_PROJECT_ID) update_method.assert_called_once_with( jobName=TEST_TRANSFER_JOB_NAME, body={ PROJECT_ID: TEST_PROJECT_ID, TRANSFER_JOB: { STATUS: GcpTransferJobsStatus.DELETED }, TRANSFER_JOB_FIELD_MASK: STATUS, }, ) execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_cancel_transfer_operation(self, get_conn): cancel_method = get_conn.return_value.transferOperations.return_value.cancel execute_method = cancel_method.return_value.execute self.gct_hook.cancel_transfer_operation( operation_name=TEST_TRANSFER_OPERATION_NAME) cancel_method.assert_called_once_with( name=TEST_TRANSFER_OPERATION_NAME) execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_get_transfer_operation(self, get_conn): get_method = get_conn.return_value.transferOperations.return_value.get execute_method = get_method.return_value.execute execute_method.return_value = TEST_TRANSFER_OPERATION res = self.gct_hook.get_transfer_operation( operation_name=TEST_TRANSFER_OPERATION_NAME) self.assertEqual(res, TEST_TRANSFER_OPERATION) get_method.assert_called_once_with(name=TEST_TRANSFER_OPERATION_NAME) execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_list_transfer_operation(self, get_conn): list_method = get_conn.return_value.transferOperations.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = { OPERATIONS: [TEST_TRANSFER_OPERATION] } list_next = get_conn.return_value.transferOperations.return_value.list_next list_next.return_value = None res = self.gct_hook.list_transfer_operations( request_filter=TEST_TRANSFER_OPERATION_FILTER) self.assertIsNotNone(res) self.assertEqual(res, [TEST_TRANSFER_OPERATION]) list_method.assert_called_once_with(filter=mock.ANY, name='transferOperations') args, kwargs = list_method.call_args_list[0] self.assertEqual( json.loads(kwargs['filter']), { FILTER_PROJECT_ID: TEST_PROJECT_ID, FILTER_JOB_NAMES: [TEST_TRANSFER_JOB_NAME] }, ) list_execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_pause_transfer_operation(self, get_conn): pause_method = get_conn.return_value.transferOperations.return_value.pause execute_method = pause_method.return_value.execute self.gct_hook.pause_transfer_operation( operation_name=TEST_TRANSFER_OPERATION_NAME) pause_method.assert_called_once_with(name=TEST_TRANSFER_OPERATION_NAME) execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_resume_transfer_operation(self, get_conn): resume_method = get_conn.return_value.transferOperations.return_value.resume execute_method = resume_method.return_value.execute self.gct_hook.resume_transfer_operation( operation_name=TEST_TRANSFER_OPERATION_NAME) resume_method.assert_called_once_with( name=TEST_TRANSFER_OPERATION_NAME) execute_method.assert_called_once_with(num_retries=5) @mock.patch('time.sleep') @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_wait_for_transfer_job(self, get_conn, mock_sleep): list_method = get_conn.return_value.transferOperations.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.side_effect = [ { OPERATIONS: [{ METADATA: { STATUS: GcpTransferOperationStatus.IN_PROGRESS } }] }, { OPERATIONS: [{ METADATA: { STATUS: GcpTransferOperationStatus.SUCCESS } }] }, ] get_conn.return_value.transferOperations.return_value.list_next.return_value = None self.gct_hook.wait_for_transfer_job({ PROJECT_ID: TEST_PROJECT_ID, 'name': 'transferJobs/test-job' }) self.assertTrue(list_method.called) list_method.assert_called_with(name='transferOperations', filter=mock.ANY) args, kwargs = list_method.call_args_list[0] self.assertEqual( json.loads(kwargs['filter']), { FILTER_PROJECT_ID: TEST_PROJECT_ID, FILTER_JOB_NAMES: ["transferJobs/test-job"] }, ) mock_sleep.assert_called_once_with(TIME_TO_SLEEP_IN_SECONDS) @mock.patch('time.sleep') @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_wait_for_transfer_job_failed(self, get_conn, mock_sleep): # pylint: disable=unused-argument list_method = get_conn.return_value.transferOperations.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = { OPERATIONS: [{ NAME: TEST_TRANSFER_OPERATION_NAME, METADATA: { STATUS: GcpTransferOperationStatus.FAILED } }] } get_conn.return_value.transferOperations.return_value.list_next.return_value = None with self.assertRaises(AirflowException): self.gct_hook.wait_for_transfer_job({ PROJECT_ID: TEST_PROJECT_ID, NAME: 'transferJobs/test-job' }) self.assertTrue(list_method.called) @mock.patch('time.sleep') @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_wait_for_transfer_job_expect_failed(self, get_conn, mock_sleep): # pylint: disable=unused-argument list_method = get_conn.return_value.transferOperations.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = { OPERATIONS: [{ NAME: TEST_TRANSFER_OPERATION_NAME, METADATA: { STATUS: GcpTransferOperationStatus.FAILED } }] } get_conn.return_value.transferOperations.return_value.list_next.return_value = None with six.assertRaisesRegex( self, AirflowException, "An unexpected operation status was encountered. Expected: SUCCESS" ): self.gct_hook.wait_for_transfer_job( job={ PROJECT_ID: 'test-project', NAME: 'transferJobs/test-job' }, expected_statuses=GcpTransferOperationStatus.SUCCESS, ) @parameterized.expand([ ([GcpTransferOperationStatus.ABORTED], (GcpTransferOperationStatus.IN_PROGRESS, )), ( [ GcpTransferOperationStatus.SUCCESS, GcpTransferOperationStatus.ABORTED ], (GcpTransferOperationStatus.IN_PROGRESS, ), ), ( [ GcpTransferOperationStatus.PAUSED, GcpTransferOperationStatus.ABORTED ], (GcpTransferOperationStatus.IN_PROGRESS, ), ), ([GcpTransferOperationStatus.ABORTED], (GcpTransferOperationStatus.IN_PROGRESS, )), ( [ GcpTransferOperationStatus.SUCCESS, GcpTransferOperationStatus.ABORTED ], (GcpTransferOperationStatus.IN_PROGRESS, ), ), ( [ GcpTransferOperationStatus.PAUSED, GcpTransferOperationStatus.ABORTED ], (GcpTransferOperationStatus.IN_PROGRESS, ), ), ]) 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 six.assertRaisesRegex( self, AirflowException, "An unexpected operation status was encountered. Expected: {}". format(", ".join(expected_statuses)), ): GCPTransferServiceHook.operations_contain_expected_statuses( operations, GcpTransferOperationStatus.IN_PROGRESS) @parameterized.expand([ ([GcpTransferOperationStatus.ABORTED], GcpTransferOperationStatus.ABORTED), ( [ GcpTransferOperationStatus.SUCCESS, GcpTransferOperationStatus.ABORTED ], GcpTransferOperationStatus.ABORTED, ), ( [ GcpTransferOperationStatus.PAUSED, GcpTransferOperationStatus.ABORTED ], GcpTransferOperationStatus.ABORTED, ), ([GcpTransferOperationStatus.ABORTED], (GcpTransferOperationStatus.ABORTED, )), ( [ GcpTransferOperationStatus.SUCCESS, GcpTransferOperationStatus.ABORTED ], (GcpTransferOperationStatus.ABORTED, ), ), ( [ GcpTransferOperationStatus.PAUSED, GcpTransferOperationStatus.ABORTED ], (GcpTransferOperationStatus.ABORTED, ), ), ]) 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)
class TestGCPTransferServiceHookWithoutProjectId(unittest.TestCase): def setUp(self): with mock.patch( 'airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__', new=mock_base_gcp_hook_no_default_project_id, ): self.gct_hook = GCPTransferServiceHook(gcp_conn_id='test') @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_create_transfer_job(self, get_conn): create_method = get_conn.return_value.transferJobs.return_value.create execute_method = create_method.return_value.execute execute_method.return_value = deepcopy(TEST_TRANSFER_JOB) with self.assertRaises(AirflowException) as e: self.gct_hook.create_transfer_job( body=_without_key(TEST_BODY, PROJECT_ID)) self.assertEqual( 'The project id must be passed either as `projectId` key in `body` parameter or as project_id ' 'extra in GCP connection definition. Both are not set!', str(e.exception), ) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_get_transfer_job(self, get_conn): get_method = get_conn.return_value.transferJobs.return_value.get execute_method = get_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB with self.assertRaises(AirflowException) as e: self.gct_hook.get_transfer_job(job_name=TEST_TRANSFER_JOB_NAME) self.assertEqual( 'The project id must be passed either as keyword project_id ' 'parameter or as project_id extra in GCP connection definition. ' 'Both are not set!', str(e.exception), ) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_list_transfer_job(self, get_conn): list_method = get_conn.return_value.transferJobs.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = { "transferJobs": [TEST_TRANSFER_JOB] } list_next = get_conn.return_value.transferJobs.return_value.list_next list_next.return_value = None with self.assertRaises(AirflowException) as e: self.gct_hook.list_transfer_job(request_filter=_without_key( TEST_TRANSFER_JOB_FILTER, FILTER_PROJECT_ID)) self.assertEqual( 'The project id must be passed either as `project_id` key in `filter` parameter or as ' 'project_id extra in GCP connection definition. Both are not set!', str(e.exception), ) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_list_transfer_operation_multiple_page(self, get_conn): pages_requests = [ mock.Mock(**{ 'execute.return_value': { "operations": [TEST_TRANSFER_OPERATION] } }) for _ in range(4) ] transfer_operation_mock = mock.Mock( **{ 'list.return_value': pages_requests[1], 'list_next.side_effect': pages_requests[1:] + [None] }) get_conn.return_value.transferOperations.return_value = transfer_operation_mock res = self.gct_hook.list_transfer_operations( request_filter=TEST_TRANSFER_OPERATION_FILTER) self.assertEqual(res, [TEST_TRANSFER_OPERATION] * 4) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_update_transfer_job(self, get_conn): update_method = get_conn.return_value.transferJobs.return_value.patch execute_method = update_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB with self.assertRaises(AirflowException) as e: self.gct_hook.update_transfer_job( job_name=TEST_TRANSFER_JOB_NAME, body=_without_key(TEST_UPDATE_TRANSFER_JOB_BODY, PROJECT_ID)) self.assertEqual( 'The project id must be passed either as `projectId` key in `body` parameter or as project_id ' 'extra in GCP connection definition. Both are not set!', str(e.exception), ) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_delete_transfer_job(self, get_conn): # pylint: disable=unused-argument with self.assertRaises(AirflowException) as e: self.gct_hook.delete_transfer_job( # pylint: disable=no-value-for-parameter job_name=TEST_TRANSFER_JOB_NAME) self.assertEqual( 'The project id must be passed either as keyword project_id parameter or as project_id extra in ' 'GCP connection definition. Both are not set!', str(e.exception), ) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_list_transfer_operation(self, get_conn): list_method = get_conn.return_value.transferOperations.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = { "operations": [TEST_TRANSFER_OPERATION] } list_next = get_conn.return_value.transferOperations.return_value.list_next list_next.return_value = None with self.assertRaises(AirflowException) as e: self.gct_hook.list_transfer_operations(request_filter=_without_key( TEST_TRANSFER_OPERATION_FILTER, FILTER_PROJECT_ID)) self.assertEqual( 'The project id must be passed either as `project_id` key in `filter` parameter or as project_id ' 'extra in GCP connection definition. Both are not set!', str(e.exception), )
class TestGCPTransferServiceHookWithProjectIdFromConnection(unittest.TestCase): def setUp(self): with mock.patch( 'airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__', new=mock_base_gcp_hook_default_project_id, ): self.gct_hook = GCPTransferServiceHook(gcp_conn_id='test') @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_create_transfer_job(self, get_conn): create_method = get_conn.return_value.transferJobs.return_value.create execute_method = create_method.return_value.execute execute_method.return_value = deepcopy(TEST_TRANSFER_JOB) res = self.gct_hook.create_transfer_job( body=self._without_project_id(TEST_BODY)) self.assertEqual(res, TEST_TRANSFER_JOB) create_method.assert_called_once_with( body=self._with_project_id(TEST_BODY, 'example-project')) execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_get_transfer_job(self, get_conn): get_method = get_conn.return_value.transferJobs.return_value.get execute_method = get_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB res = self.gct_hook.get_transfer_job(job_name=TEST_TRANSFER_JOB_NAME) self.assertIsNotNone(res) self.assertEqual(TEST_TRANSFER_JOB_NAME, res[NAME]) get_method.assert_called_once_with(jobName=TEST_TRANSFER_JOB_NAME, projectId='example-project') execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_list_transfer_job(self, get_conn): list_method = get_conn.return_value.transferJobs.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = {TRANSFER_JOBS: [TEST_TRANSFER_JOB]} list_next = get_conn.return_value.transferJobs.return_value.list_next list_next.return_value = None res = self.gct_hook.list_transfer_job(request_filter=_without_key( TEST_TRANSFER_JOB_FILTER, FILTER_PROJECT_ID)) self.assertIsNotNone(res) self.assertEqual(res, [TEST_TRANSFER_JOB]) list_method.assert_called_with(filter=mock.ANY) args, kwargs = list_method.call_args_list[0] self.assertEqual( json.loads(kwargs['filter']), { FILTER_PROJECT_ID: 'example-project', FILTER_JOB_NAMES: [TEST_TRANSFER_JOB_NAME] }, ) list_execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_update_transfer_job(self, get_conn): update_method = get_conn.return_value.transferJobs.return_value.patch execute_method = update_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB res = self.gct_hook.update_transfer_job( job_name=TEST_TRANSFER_JOB_NAME, body=self._without_project_id(TEST_UPDATE_TRANSFER_JOB_BODY)) self.assertIsNotNone(res) update_method.assert_called_once_with( jobName=TEST_TRANSFER_JOB_NAME, body=self._with_project_id(TEST_UPDATE_TRANSFER_JOB_BODY, 'example-project'), ) execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_delete_transfer_job(self, get_conn): update_method = get_conn.return_value.transferJobs.return_value.patch execute_method = update_method.return_value.execute self.gct_hook.delete_transfer_job(job_name=TEST_TRANSFER_JOB_NAME, project_id=TEST_PROJECT_ID) update_method.assert_called_once_with( jobName=TEST_TRANSFER_JOB_NAME, body={ PROJECT_ID: TEST_PROJECT_ID, TRANSFER_JOB: { STATUS: 'DELETED' }, TRANSFER_JOB_FIELD_MASK: STATUS, }, ) execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_cancel_transfer_operation(self, get_conn): cancel_method = get_conn.return_value.transferOperations.return_value.cancel execute_method = cancel_method.return_value.execute self.gct_hook.cancel_transfer_operation( operation_name=TEST_TRANSFER_OPERATION_NAME) cancel_method.assert_called_once_with( name=TEST_TRANSFER_OPERATION_NAME) execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_get_transfer_operation(self, get_conn): get_method = get_conn.return_value.transferOperations.return_value.get execute_method = get_method.return_value.execute execute_method.return_value = TEST_TRANSFER_OPERATION res = self.gct_hook.get_transfer_operation( operation_name=TEST_TRANSFER_OPERATION_NAME) self.assertEqual(res, TEST_TRANSFER_OPERATION) get_method.assert_called_once_with(name=TEST_TRANSFER_OPERATION_NAME) execute_method.assert_called_once_with(num_retries=5) @mock.patch( 'airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn' ) def test_list_transfer_operation(self, get_conn): list_method = get_conn.return_value.transferOperations.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = { OPERATIONS: [TEST_TRANSFER_OPERATION] } list_next = get_conn.return_value.transferOperations.return_value.list_next list_next.return_value = None res = self.gct_hook.list_transfer_operations( request_filter=_without_key(TEST_TRANSFER_OPERATION_FILTER, FILTER_PROJECT_ID)) self.assertIsNotNone(res) self.assertEqual(res, [TEST_TRANSFER_OPERATION]) list_method.assert_called_once_with(filter=mock.ANY, name='transferOperations') args, kwargs = list_method.call_args_list[0] self.assertEqual( json.loads(kwargs['filter']), { FILTER_PROJECT_ID: 'example-project', FILTER_JOB_NAMES: [TEST_TRANSFER_JOB_NAME] }, ) list_execute_method.assert_called_once_with(num_retries=5) @staticmethod def _without_project_id(body): body = deepcopy(body) del body[PROJECT_ID] return body @staticmethod def _with_project_id(body, project_id): body = deepcopy(body) del body[PROJECT_ID] body[PROJECT_ID] = project_id return body
class TestGCPTransferServiceHook(unittest.TestCase): def setUp(self): with mock.patch.object(GCPTransferServiceHook, '__init__', return_value=None): self.conn = mock.Mock() self.transfer_hook = GCPTransferServiceHook() self.transfer_hook._conn = self.conn def test_create_transfer_job(self): mock_create = self.conn.transferJobs.return_value.create mock_execute = mock_create.return_value.execute mock_execute.return_value = { 'projectId': 'test-project', 'name': 'transferJobs/test-job', } now = datetime.datetime.utcnow() transfer_spec = { 'awsS3DataSource': {'bucketName': 'test-s3-bucket'}, 'gcsDataSink': {'bucketName': 'test-gcs-bucket'} } self.transfer_hook.create_transfer_job( project_id='test-project', description='test-description', schedule=None, transfer_spec=transfer_spec) mock_create.assert_called_once_with(body={ 'status': 'ENABLED', 'projectId': 'test-project', 'description': 'test-description', 'transferSpec': transfer_spec, 'schedule': { 'scheduleStartDate': { 'day': now.day, 'month': now.month, 'year': now.year, }, 'scheduleEndDate': { 'day': now.day, 'month': now.month, 'year': now.year, } } }) def test_create_transfer_job_custom_schedule(self): mock_create = self.conn.transferJobs.return_value.create mock_execute = mock_create.return_value.execute mock_execute.return_value = { 'projectId': 'test-project', 'name': 'transferJobs/test-job', } schedule = { 'scheduleStartDate': {'month': 10, 'day': 1, 'year': 2018}, 'scheduleEndDate': {'month': 10, 'day': 31, 'year': 2018}, } transfer_spec = { 'awsS3DataSource': {'bucketName': 'test-s3-bucket'}, 'gcsDataSink': {'bucketName': 'test-gcs-bucket'} } self.transfer_hook.create_transfer_job( project_id='test-project', description='test-description', schedule=schedule, transfer_spec=transfer_spec) mock_create.assert_called_once_with(body={ 'status': 'ENABLED', 'projectId': 'test-project', 'description': 'test-description', 'transferSpec': transfer_spec, 'schedule': schedule, }) @mock.patch('time.sleep') def test_wait_for_transfer_job(self, mock_sleep): mock_list = self.conn.transferOperations.return_value.list mock_execute = mock_list.return_value.execute mock_execute.side_effect = [ {'operations': [{'metadata': {'status': 'IN_PROGRESS'}}]}, {'operations': [{'metadata': {'status': 'SUCCESS'}}]}, ] self.transfer_hook.wait_for_transfer_job({ 'projectId': 'test-project', 'name': 'transferJobs/test-job', }) self.assertTrue(mock_list.called) list_args, list_kwargs = mock_list.call_args_list[0] self.assertEqual(list_kwargs.get('name'), 'transferOperations') self.assertEqual( json.loads(list_kwargs.get('filter')), { 'project_id': 'test-project', 'job_names': ['transferJobs/test-job'] }, ) mock_sleep.assert_called_once_with(TIME_TO_SLEEP_IN_SECONDS) def test_wait_for_transfer_job_failed(self): mock_list = self.conn.transferOperations.return_value.list mock_execute = mock_list.return_value.execute mock_execute.side_effect = [ {'operations': [{'name': 'test-job', 'metadata': {'status': 'FAILED'}}]}, ] with self.assertRaises(AirflowException): self.transfer_hook.wait_for_transfer_job({ 'projectId': 'test-project', 'name': 'transferJobs/test-job', })
class TestGCPTransferServiceHook(unittest.TestCase): def setUp(self): with mock.patch.object(GCPTransferServiceHook, '__init__', return_value=None): self.conn = mock.Mock() self.transfer_hook = GCPTransferServiceHook() self.transfer_hook._conn = self.conn def test_create_transfer_job(self): mock_create = self.conn.transferJobs.return_value.create mock_execute = mock_create.return_value.execute mock_execute.return_value = { 'projectId': 'test-project', 'name': 'transferJobs/test-job', } now = datetime.datetime.utcnow() transfer_spec = { 'awsS3DataSource': {'bucketName': 'test-s3-bucket'}, 'gcsDataSink': {'bucketName': 'test-gcs-bucket'} } self.transfer_hook.create_transfer_job( 'test-project', 'test-description', None, transfer_spec) mock_create.assert_called_once_with(body={ 'status': 'ENABLED', 'projectId': 'test-project', 'description': 'test-description', 'transferSpec': transfer_spec, 'schedule': { 'scheduleStartDate': { 'day': now.day, 'month': now.month, 'year': now.year, }, 'scheduleEndDate': { 'day': now.day, 'month': now.month, 'year': now.year, } } }) def test_create_transfer_job_custom_schedule(self): mock_create = self.conn.transferJobs.return_value.create mock_execute = mock_create.return_value.execute mock_execute.return_value = { 'projectId': 'test-project', 'name': 'transferJobs/test-job', } schedule = { 'scheduleStartDate': {'month': 10, 'day': 1, 'year': 2018}, 'scheduleEndDate': {'month': 10, 'day': 31, 'year': 2018}, } transfer_spec = { 'awsS3DataSource': {'bucketName': 'test-s3-bucket'}, 'gcsDataSink': {'bucketName': 'test-gcs-bucket'} } self.transfer_hook.create_transfer_job( 'test-project', 'test-description', schedule, transfer_spec) mock_create.assert_called_once_with(body={ 'status': 'ENABLED', 'projectId': 'test-project', 'description': 'test-description', 'transferSpec': transfer_spec, 'schedule': schedule, }) @mock.patch('time.sleep') def test_wait_for_transfer_job(self, mock_sleep): mock_list = self.conn.transferOperations.return_value.list mock_execute = mock_list.return_value.execute mock_execute.side_effect = [ {'operations': [{'metadata': {'status': 'IN_PROGRESS'}}]}, {'operations': [{'metadata': {'status': 'SUCCESS'}}]}, ] self.transfer_hook.wait_for_transfer_job({ 'projectId': 'test-project', 'name': 'transferJobs/test-job', }) self.assertTrue(mock_list.called) list_args, list_kwargs = mock_list.call_args_list[0] self.assertEqual(list_kwargs.get('name'), 'transferOperations') self.assertEqual( json.loads(list_kwargs.get('filter')), { 'project_id': 'test-project', 'job_names': ['transferJobs/test-job'] }, ) mock_sleep.assert_called_once_with(TIME_TO_SLEEP_IN_SECONDS) def test_wait_for_transfer_job_failed(self): mock_list = self.conn.transferOperations.return_value.list mock_execute = mock_list.return_value.execute mock_execute.side_effect = [ {'operations': [{'name': 'test-job', 'metadata': {'status': 'FAILED'}}]}, ] with self.assertRaises(AirflowException): self.transfer_hook.wait_for_transfer_job({ 'projectId': 'test-project', 'name': 'transferJobs/test-job', })
class TestGCPTransferServiceHookWithPassedProjectId(unittest.TestCase): def setUp(self): with mock.patch( 'airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__', new=mock_base_gcp_hook_no_default_project_id, ): self.gct_hook = GCPTransferServiceHook(gcp_conn_id='test') @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_create_transfer_job(self, get_conn): create_method = get_conn.return_value.transferJobs.return_value.create execute_method = create_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB res = self.gct_hook.create_transfer_job(body=TEST_BODY) self.assertEqual(res, TEST_TRANSFER_JOB) create_method.assert_called_once_with(body=TEST_BODY) execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_get_transfer_job(self, get_conn): get_method = get_conn.return_value.transferJobs.return_value.get execute_method = get_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB res = self.gct_hook.get_transfer_job(job_name=TEST_TRANSFER_JOB_NAME, project_id=TEST_PROJECT_ID) self.assertIsNotNone(res) self.assertEqual(TEST_TRANSFER_JOB_NAME, res[NAME]) get_method.assert_called_once_with(jobName=TEST_TRANSFER_JOB_NAME, projectId=TEST_PROJECT_ID) execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_list_transfer_job(self, get_conn): list_method = get_conn.return_value.transferJobs.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = {TRANSFER_JOBS: [TEST_TRANSFER_JOB]} list_next = get_conn.return_value.transferJobs.return_value.list_next list_next.return_value = None res = self.gct_hook.list_transfer_job(filter=TEST_TRANSFER_JOB_FILTER) self.assertIsNotNone(res) self.assertEqual(res, [TEST_TRANSFER_JOB]) list_method.assert_called_once_with(filter=mock.ANY) args, kwargs = list_method.call_args_list[0] self.assertEqual( json.loads(kwargs['filter']), {FILTER_PROJECT_ID: TEST_PROJECT_ID, FILTER_JOB_NAMES: [TEST_TRANSFER_JOB_NAME]}, ) list_execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_update_transfer_job(self, get_conn): update_method = get_conn.return_value.transferJobs.return_value.patch execute_method = update_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB res = self.gct_hook.update_transfer_job( job_name=TEST_TRANSFER_JOB_NAME, body=TEST_UPDATE_TRANSFER_JOB_BODY ) self.assertIsNotNone(res) update_method.assert_called_once_with( jobName=TEST_TRANSFER_JOB_NAME, body=TEST_UPDATE_TRANSFER_JOB_BODY ) execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_delete_transfer_job(self, get_conn): update_method = get_conn.return_value.transferJobs.return_value.patch execute_method = update_method.return_value.execute self.gct_hook.delete_transfer_job(job_name=TEST_TRANSFER_JOB_NAME, project_id=TEST_PROJECT_ID) update_method.assert_called_once_with( jobName=TEST_TRANSFER_JOB_NAME, body={ PROJECT_ID: TEST_PROJECT_ID, TRANSFER_JOB: {STATUS: GcpTransferJobsStatus.DELETED}, TRANSFER_JOB_FIELD_MASK: STATUS, }, ) execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_cancel_transfer_operation(self, get_conn): cancel_method = get_conn.return_value.transferOperations.return_value.cancel execute_method = cancel_method.return_value.execute self.gct_hook.cancel_transfer_operation(operation_name=TEST_TRANSFER_OPERATION_NAME) cancel_method.assert_called_once_with(name=TEST_TRANSFER_OPERATION_NAME) execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_get_transfer_operation(self, get_conn): get_method = get_conn.return_value.transferOperations.return_value.get execute_method = get_method.return_value.execute execute_method.return_value = TEST_TRANSFER_OPERATION res = self.gct_hook.get_transfer_operation(operation_name=TEST_TRANSFER_OPERATION_NAME) self.assertEqual(res, TEST_TRANSFER_OPERATION) get_method.assert_called_once_with(name=TEST_TRANSFER_OPERATION_NAME) execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_list_transfer_operation(self, get_conn): list_method = get_conn.return_value.transferOperations.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = {OPERATIONS: [TEST_TRANSFER_OPERATION]} list_next = get_conn.return_value.transferOperations.return_value.list_next list_next.return_value = None res = self.gct_hook.list_transfer_operations(filter=TEST_TRANSFER_OPERATION_FILTER) self.assertIsNotNone(res) self.assertEqual(res, [TEST_TRANSFER_OPERATION]) list_method.assert_called_once_with(filter=mock.ANY, name='transferOperations') args, kwargs = list_method.call_args_list[0] self.assertEqual( json.loads(kwargs['filter']), {FILTER_PROJECT_ID: TEST_PROJECT_ID, FILTER_JOB_NAMES: [TEST_TRANSFER_JOB_NAME]}, ) list_execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_pause_transfer_operation(self, get_conn): pause_method = get_conn.return_value.transferOperations.return_value.pause execute_method = pause_method.return_value.execute self.gct_hook.pause_transfer_operation(operation_name=TEST_TRANSFER_OPERATION_NAME) pause_method.assert_called_once_with(name=TEST_TRANSFER_OPERATION_NAME) execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_resume_transfer_operation(self, get_conn): resume_method = get_conn.return_value.transferOperations.return_value.resume execute_method = resume_method.return_value.execute self.gct_hook.resume_transfer_operation(operation_name=TEST_TRANSFER_OPERATION_NAME) resume_method.assert_called_once_with(name=TEST_TRANSFER_OPERATION_NAME) execute_method.assert_called_once_with(num_retries=5) @mock.patch('time.sleep') @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_wait_for_transfer_job(self, get_conn, mock_sleep): list_method = get_conn.return_value.transferOperations.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.side_effect = [ {OPERATIONS: [{METADATA: {STATUS: GcpTransferOperationStatus.IN_PROGRESS}}]}, {OPERATIONS: [{METADATA: {STATUS: GcpTransferOperationStatus.SUCCESS}}]}, ] get_conn.return_value.transferOperations.return_value.list_next.return_value = None self.gct_hook.wait_for_transfer_job({PROJECT_ID: TEST_PROJECT_ID, 'name': 'transferJobs/test-job'}) self.assertTrue(list_method.called) list_method.assert_called_with(name='transferOperations', filter=mock.ANY) args, kwargs = list_method.call_args_list[0] self.assertEqual( json.loads(kwargs['filter']), {FILTER_PROJECT_ID: TEST_PROJECT_ID, FILTER_JOB_NAMES: ["transferJobs/test-job"]}, ) mock_sleep.assert_called_once_with(TIME_TO_SLEEP_IN_SECONDS) @mock.patch('time.sleep') @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_wait_for_transfer_job_failed(self, get_conn, mock_sleep): list_method = get_conn.return_value.transferOperations.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = { OPERATIONS: [ {NAME: TEST_TRANSFER_OPERATION_NAME, METADATA: {STATUS: GcpTransferOperationStatus.FAILED}} ] } get_conn.return_value.transferOperations.return_value.list_next.return_value = None with self.assertRaises(AirflowException): self.gct_hook.wait_for_transfer_job({PROJECT_ID: TEST_PROJECT_ID, NAME: 'transferJobs/test-job'}) self.assertTrue(list_method.called) @mock.patch('time.sleep') @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_wait_for_transfer_job_expect_failed(self, get_conn, mock_sleep): list_method = get_conn.return_value.transferOperations.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = { OPERATIONS: [ {NAME: TEST_TRANSFER_OPERATION_NAME, METADATA: {STATUS: GcpTransferOperationStatus.FAILED}} ] } get_conn.return_value.transferOperations.return_value.list_next.return_value = None with six.assertRaisesRegex( self, AirflowException, "An unexpected operation status was encountered. Expected: SUCCESS" ): self.gct_hook.wait_for_transfer_job( job={PROJECT_ID: 'test-project', NAME: 'transferJobs/test-job'}, expected_statuses=GcpTransferOperationStatus.SUCCESS, ) @parameterized.expand( [ ([GcpTransferOperationStatus.ABORTED], (GcpTransferOperationStatus.IN_PROGRESS,)), ( [GcpTransferOperationStatus.SUCCESS, GcpTransferOperationStatus.ABORTED], (GcpTransferOperationStatus.IN_PROGRESS,), ), ( [GcpTransferOperationStatus.PAUSED, GcpTransferOperationStatus.ABORTED], (GcpTransferOperationStatus.IN_PROGRESS,), ), ([GcpTransferOperationStatus.ABORTED], (GcpTransferOperationStatus.IN_PROGRESS,)), ( [GcpTransferOperationStatus.SUCCESS, GcpTransferOperationStatus.ABORTED], (GcpTransferOperationStatus.IN_PROGRESS,), ), ( [GcpTransferOperationStatus.PAUSED, GcpTransferOperationStatus.ABORTED], (GcpTransferOperationStatus.IN_PROGRESS,), ), ] ) 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 six.assertRaisesRegex( self, AirflowException, "An unexpected operation status was encountered. Expected: {}".format( ", ".join(expected_statuses) ), ): GCPTransferServiceHook.operations_contain_expected_statuses( operations, GcpTransferOperationStatus.IN_PROGRESS ) @parameterized.expand( [ ([GcpTransferOperationStatus.ABORTED], GcpTransferOperationStatus.ABORTED), ( [GcpTransferOperationStatus.SUCCESS, GcpTransferOperationStatus.ABORTED], GcpTransferOperationStatus.ABORTED, ), ( [GcpTransferOperationStatus.PAUSED, GcpTransferOperationStatus.ABORTED], GcpTransferOperationStatus.ABORTED, ), ([GcpTransferOperationStatus.ABORTED], (GcpTransferOperationStatus.ABORTED,)), ( [GcpTransferOperationStatus.SUCCESS, GcpTransferOperationStatus.ABORTED], (GcpTransferOperationStatus.ABORTED,), ), ( [GcpTransferOperationStatus.PAUSED, GcpTransferOperationStatus.ABORTED], (GcpTransferOperationStatus.ABORTED,), ), ] ) 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)
class TestGCPTransferServiceHookWithoutProjectId(unittest.TestCase): def setUp(self): with mock.patch( 'airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__', new=mock_base_gcp_hook_no_default_project_id, ): self.gct_hook = GCPTransferServiceHook(gcp_conn_id='test') @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_create_transfer_job(self, get_conn): create_method = get_conn.return_value.transferJobs.return_value.create execute_method = create_method.return_value.execute execute_method.return_value = deepcopy(TEST_TRANSFER_JOB) with self.assertRaises(AirflowException) as e: self.gct_hook.create_transfer_job(body=_without_key(TEST_BODY, PROJECT_ID)) self.assertEqual( 'The project id must be passed either as `projectId` key in `body` parameter or as project_id ' 'extra in GCP connection definition. Both are not set!', str(e.exception), ) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_get_transfer_job(self, get_conn): get_method = get_conn.return_value.transferJobs.return_value.get execute_method = get_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB with self.assertRaises(AirflowException) as e: self.gct_hook.get_transfer_job(job_name=TEST_TRANSFER_JOB_NAME) self.assertEqual( 'The project id must be passed either as keyword project_id ' 'parameter or as project_id extra in GCP connection definition. ' 'Both are not set!', str(e.exception), ) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_list_transfer_job(self, get_conn): list_method = get_conn.return_value.transferJobs.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = {"transferJobs": [TEST_TRANSFER_JOB]} list_next = get_conn.return_value.transferJobs.return_value.list_next list_next.return_value = None with self.assertRaises(AirflowException) as e: self.gct_hook.list_transfer_job(filter=_without_key(TEST_TRANSFER_JOB_FILTER, FILTER_PROJECT_ID)) self.assertEqual( 'The project id must be passed either as `project_id` key in `filter` parameter or as ' 'project_id extra in GCP connection definition. Both are not set!', str(e.exception), ) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_list_transfer_operation_multiple_page(self, get_conn): pages_requests = [ mock.Mock(**{'execute.return_value': {"operations": [TEST_TRANSFER_OPERATION]}}) for _ in range(4) ] transfer_operation_mock = mock.Mock( **{'list.return_value': pages_requests[1], 'list_next.side_effect': pages_requests[1:] + [None]} ) get_conn.return_value.transferOperations.return_value = transfer_operation_mock res = self.gct_hook.list_transfer_operations(filter=TEST_TRANSFER_OPERATION_FILTER) self.assertEqual(res, [TEST_TRANSFER_OPERATION] * 4) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_update_transfer_job(self, get_conn): update_method = get_conn.return_value.transferJobs.return_value.patch execute_method = update_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB with self.assertRaises(AirflowException) as e: self.gct_hook.update_transfer_job( job_name=TEST_TRANSFER_JOB_NAME, body=_without_key(TEST_UPDATE_TRANSFER_JOB_BODY, PROJECT_ID) ) self.assertEqual( 'The project id must be passed either as `projectId` key in `body` parameter or as project_id ' 'extra in GCP connection definition. Both are not set!', str(e.exception), ) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_delete_transfer_job(self, get_conn): with self.assertRaises(AirflowException) as e: self.gct_hook.delete_transfer_job(job_name=TEST_TRANSFER_JOB_NAME) self.assertEqual( 'The project id must be passed either as keyword project_id parameter or as project_id extra in ' 'GCP connection definition. Both are not set!', str(e.exception), ) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_list_transfer_operation(self, get_conn): list_method = get_conn.return_value.transferOperations.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = {"operations": [TEST_TRANSFER_OPERATION]} list_next = get_conn.return_value.transferOperations.return_value.list_next list_next.return_value = None with self.assertRaises(AirflowException) as e: self.gct_hook.list_transfer_operations( filter=_without_key(TEST_TRANSFER_OPERATION_FILTER, FILTER_PROJECT_ID) ) self.assertEqual( 'The project id must be passed either as `project_id` key in `filter` parameter or as project_id ' 'extra in GCP connection definition. Both are not set!', str(e.exception), )
class TestGCPTransferServiceHookWithProjectIdFromConnection(unittest.TestCase): def setUp(self): with mock.patch( 'airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__', new=mock_base_gcp_hook_default_project_id, ): self.gct_hook = GCPTransferServiceHook(gcp_conn_id='test') @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_create_transfer_job(self, get_conn): create_method = get_conn.return_value.transferJobs.return_value.create execute_method = create_method.return_value.execute execute_method.return_value = deepcopy(TEST_TRANSFER_JOB) res = self.gct_hook.create_transfer_job(body=self._without_project_id(TEST_BODY)) self.assertEqual(res, TEST_TRANSFER_JOB) create_method.assert_called_once_with(body=self._with_project_id(TEST_BODY, 'example-project')) execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_get_transfer_job(self, get_conn): get_method = get_conn.return_value.transferJobs.return_value.get execute_method = get_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB res = self.gct_hook.get_transfer_job(job_name=TEST_TRANSFER_JOB_NAME) self.assertIsNotNone(res) self.assertEqual(TEST_TRANSFER_JOB_NAME, res[NAME]) get_method.assert_called_once_with(jobName=TEST_TRANSFER_JOB_NAME, projectId='example-project') execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_list_transfer_job(self, get_conn): list_method = get_conn.return_value.transferJobs.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = {TRANSFER_JOBS: [TEST_TRANSFER_JOB]} list_next = get_conn.return_value.transferJobs.return_value.list_next list_next.return_value = None res = self.gct_hook.list_transfer_job( filter=_without_key(TEST_TRANSFER_JOB_FILTER, FILTER_PROJECT_ID) ) self.assertIsNotNone(res) self.assertEqual(res, [TEST_TRANSFER_JOB]) list_method.assert_called_with(filter=mock.ANY) args, kwargs = list_method.call_args_list[0] self.assertEqual( json.loads(kwargs['filter']), {FILTER_PROJECT_ID: 'example-project', FILTER_JOB_NAMES: [TEST_TRANSFER_JOB_NAME]}, ) list_execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_update_transfer_job(self, get_conn): update_method = get_conn.return_value.transferJobs.return_value.patch execute_method = update_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB res = self.gct_hook.update_transfer_job( job_name=TEST_TRANSFER_JOB_NAME, body=self._without_project_id(TEST_UPDATE_TRANSFER_JOB_BODY) ) self.assertIsNotNone(res) update_method.assert_called_once_with( jobName=TEST_TRANSFER_JOB_NAME, body=self._with_project_id(TEST_UPDATE_TRANSFER_JOB_BODY, 'example-project'), ) execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_delete_transfer_job(self, get_conn): update_method = get_conn.return_value.transferJobs.return_value.patch execute_method = update_method.return_value.execute self.gct_hook.delete_transfer_job(job_name=TEST_TRANSFER_JOB_NAME, project_id=TEST_PROJECT_ID) update_method.assert_called_once_with( jobName=TEST_TRANSFER_JOB_NAME, body={ PROJECT_ID: TEST_PROJECT_ID, TRANSFER_JOB: {STATUS: 'DELETED'}, TRANSFER_JOB_FIELD_MASK: STATUS, }, ) execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_cancel_transfer_operation(self, get_conn): cancel_method = get_conn.return_value.transferOperations.return_value.cancel execute_method = cancel_method.return_value.execute self.gct_hook.cancel_transfer_operation(operation_name=TEST_TRANSFER_OPERATION_NAME) cancel_method.assert_called_once_with(name=TEST_TRANSFER_OPERATION_NAME) execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_get_transfer_operation(self, get_conn): get_method = get_conn.return_value.transferOperations.return_value.get execute_method = get_method.return_value.execute execute_method.return_value = TEST_TRANSFER_OPERATION res = self.gct_hook.get_transfer_operation(operation_name=TEST_TRANSFER_OPERATION_NAME) self.assertEqual(res, TEST_TRANSFER_OPERATION) get_method.assert_called_once_with(name=TEST_TRANSFER_OPERATION_NAME) execute_method.assert_called_once_with(num_retries=5) @mock.patch('airflow.contrib.hooks.gcp_transfer_hook.GCPTransferServiceHook.get_conn') def test_list_transfer_operation(self, get_conn): list_method = get_conn.return_value.transferOperations.return_value.list list_execute_method = list_method.return_value.execute list_execute_method.return_value = {OPERATIONS: [TEST_TRANSFER_OPERATION]} list_next = get_conn.return_value.transferOperations.return_value.list_next list_next.return_value = None res = self.gct_hook.list_transfer_operations( filter=_without_key(TEST_TRANSFER_OPERATION_FILTER, FILTER_PROJECT_ID) ) self.assertIsNotNone(res) self.assertEqual(res, [TEST_TRANSFER_OPERATION]) list_method.assert_called_once_with(filter=mock.ANY, name='transferOperations') args, kwargs = list_method.call_args_list[0] self.assertEqual( json.loads(kwargs['filter']), {FILTER_PROJECT_ID: 'example-project', FILTER_JOB_NAMES: [TEST_TRANSFER_JOB_NAME]}, ) list_execute_method.assert_called_once_with(num_retries=5) @staticmethod def _without_project_id(body): body = deepcopy(body) del body[PROJECT_ID] return body @staticmethod def _with_project_id(body, project_id): body = deepcopy(body) del body[PROJECT_ID] body[PROJECT_ID] = project_id return body