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 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 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()
job = transfer_hook.create_transfer_job(
project_id=self.project_id,
description=self.description,
schedule=self.schedule,
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,
})
if self.wait:
transfer_hook.wait_for_transfer_job(job)
def execute(self, context):
hook = GCPTransferServiceHook(api_version=self.api_version,
gcp_conn_id=self.gcp_conn_id)
operations_list = hook.list_transfer_operations(
request_filter=self.filter)
self.log.info(operations_list)
return operations_list
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):
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 poke(self, context):
hook = GCPTransferServiceHook(gcp_conn_id=self.gcp_cloud_conn_id)
operations = hook.list_transfer_operations(
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
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
)
def poke(self, context):
hook = GCPTransferServiceHook(gcp_conn_id=self.gcp_cloud_conn_id)
operations = hook.list_transfer_operations(
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
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 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 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 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
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
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.update_transfer_job(job_name=self.job_name, 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)
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 execute(self, context):
hook = GCPTransferServiceHook(api_version=self.api_version, gcp_conn_id=self.gcp_conn_id)
hook.cancel_transfer_operation(operation_name=self.operation_name)
def execute(self, context):
hook = GCPTransferServiceHook(api_version=self.api_version, gcp_conn_id=self.gcp_conn_id)
operations_list = hook.list_transfer_operations(filter=self.filter)
self.log.info(operations_list)
return operations_list
def execute(self, context):
self._validate_inputs()
hook = GCPTransferServiceHook(api_version=self.api_version, gcp_conn_id=self.gcp_conn_id)
hook.delete_transfer_job(job_name=self.job_name, project_id=self.project_id)
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.update_transfer_job(job_name=self.job_name, body=self.body)
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
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')
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 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)
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')
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',
})
def execute(self, context):
self._validate_inputs()
hook = GCPTransferServiceHook(api_version=self.api_version,
gcp_conn_id=self.gcp_conn_id)
hook.delete_transfer_job(job_name=self.job_name,
project_id=self.project_id)
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),
)
def execute(self, context):
hook = GCPTransferServiceHook(api_version=self.api_version,
gcp_conn_id=self.gcp_conn_id)
hook.cancel_transfer_operation(operation_name=self.operation_name)
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',
})
