def test_execute_move_single_file(self, sftp_hook, gcs_hook):
task = GoogleCloudStorageToSFTPOperator(
task_id=TASK_ID,
source_bucket=TEST_BUCKET,
source_object=SOURCE_OBJECT_NO_WILDCARD,
destination_path=DESTINATION_SFTP,
move_object=True,
gcp_conn_id=GCP_CONN_ID,
sftp_conn_id=SFTP_CONN_ID,
delegate_to=DELEGATE_TO,
)
task.execute(None)
gcs_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID, delegate_to=DELEGATE_TO
)
sftp_hook.assert_called_once_with(SFTP_CONN_ID)
args, kwargs = gcs_hook.return_value.download.call_args
self.assertEqual(kwargs["bucket_name"], TEST_BUCKET)
self.assertEqual(kwargs["object_name"], SOURCE_OBJECT_NO_WILDCARD)
args, kwargs = sftp_hook.return_value.store_file.call_args
self.assertEqual(
args[0], os.path.join(DESTINATION_SFTP, SOURCE_OBJECT_NO_WILDCARD)
)
gcs_hook.return_value.delete.assert_called_once_with(
TEST_BUCKET, SOURCE_OBJECT_NO_WILDCARD
)
def test_execute_more_than_one_wildcard_exception(self, sftp_hook, gcs_hook):
gcs_hook.return_value.list.return_value = SOURCE_FILES_LIST[:2]
operator = GoogleCloudStorageToSFTPOperator(
task_id=TASK_ID,
source_bucket=TEST_BUCKET,
source_object=SOURCE_OBJECT_MULTIPLE_WILDCARDS,
destination_path=DESTINATION_SFTP,
move_object=False,
gcp_conn_id=GCP_CONN_ID,
sftp_conn_id=SFTP_CONN_ID,
delegate_to=DELEGATE_TO,
)
with self.assertRaises(AirflowException):
operator.execute(None)
def test_execute_move_with_wildcard(self, sftp_hook, gcs_hook):
gcs_hook.return_value.list.return_value = SOURCE_FILES_LIST[:2]
operator = GoogleCloudStorageToSFTPOperator(
task_id=TASK_ID,
source_bucket=TEST_BUCKET,
source_object=SOURCE_OBJECT_WILDCARD_FILENAME,
destination_path=DESTINATION_SFTP,
move_object=True,
gcp_conn_id=GCP_CONN_ID,
sftp_conn_id=SFTP_CONN_ID,
delegate_to=DELEGATE_TO,
)
operator.execute(None)
gcs_hook.return_value.list.assert_called_with(
TEST_BUCKET, delimiter=".txt", prefix="test_object"
)
call_one, call_two = gcs_hook.return_value.delete.call_args_list
self.assertEqual(call_one[0], (TEST_BUCKET, "test_object/file1.txt"))
self.assertEqual(call_two[0], (TEST_BUCKET, "test_object/file2.txt"))