def test_execute_move_single_file(self, sftp_hook, gcs_hook):
task = SFTPToGCSOperator(
task_id=TASK_ID,
source_path=SOURCE_OBJECT_NO_WILDCARD,
destination_bucket=TEST_BUCKET,
destination_path=DESTINATION_PATH_FILE,
move_object=True,
gcp_conn_id=GCP_CONN_ID,
sftp_conn_id=SFTP_CONN_ID,
delegate_to=DELEGATE_TO,
impersonation_chain=IMPERSONATION_CHAIN,
)
task.execute(None)
gcs_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
delegate_to=DELEGATE_TO,
impersonation_chain=IMPERSONATION_CHAIN,
)
sftp_hook.assert_called_once_with(SFTP_CONN_ID)
sftp_hook.return_value.retrieve_file.assert_called_once_with(
os.path.join(SOURCE_OBJECT_NO_WILDCARD), mock.ANY
)
gcs_hook.return_value.upload.assert_called_once_with(
bucket_name=TEST_BUCKET,
object_name=DESTINATION_PATH_FILE,
filename=mock.ANY,
mime_type=DEFAULT_MIME_TYPE,
)
sftp_hook.return_value.delete_file.assert_called_once_with(SOURCE_OBJECT_NO_WILDCARD)
def test_execute_move_with_wildcard(self, sftp_hook, gcs_hook):
sftp_hook.return_value.get_tree_map.return_value = [
["main_dir/test_object3.json", "main_dir/sub_dir/test_object3.json"],
[],
[],
]
gcs_hook.return_value.list.return_value = SOURCE_FILES_LIST[:2]
task = SFTPToGCSOperator(
task_id=TASK_ID,
source_path=SOURCE_OBJECT_WILDCARD_FILENAME,
destination_bucket=TEST_BUCKET,
destination_path=DESTINATION_PATH_DIR,
move_object=True,
gcp_conn_id=GCP_CONN_ID,
sftp_conn_id=SFTP_CONN_ID,
delegate_to=DELEGATE_TO,
)
task.execute(None)
sftp_hook.return_value.delete_file.assert_has_calls(
[
mock.call("main_dir/test_object3.json"),
mock.call("main_dir/sub_dir/test_object3.json"),
]
)
def test_execute_more_than_one_wildcard_exception(self, sftp_hook, gcs_hook):
task = SFTPToGCSOperator(
task_id=TASK_ID,
source_path=SOURCE_OBJECT_MULTIPLE_WILDCARDS,
destination_bucket=TEST_BUCKET,
destination_path=DESTINATION_PATH_FILE,
move_object=False,
gcp_conn_id=GCP_CONN_ID,
sftp_conn_id=SFTP_CONN_ID,
delegate_to=DELEGATE_TO,
)
with self.assertRaises(AirflowException) as cm:
task.execute(None)
err = cm.exception
self.assertIn("Only one wildcard '*' is allowed in source_path parameter", str(err))
def test_execute_copy_with_wildcard(self, sftp_hook, gcs_hook):
sftp_hook.return_value.get_tree_map.return_value = [
["main_dir/test_object3.json", "main_dir/sub_dir/test_object3.json"],
[],
[],
]
task = SFTPToGCSOperator(
task_id=TASK_ID,
source_path=SOURCE_OBJECT_WILDCARD_FILENAME,
destination_bucket=TEST_BUCKET,
destination_path=DESTINATION_PATH_DIR,
move_object=True,
gcp_conn_id=GCP_CONN_ID,
sftp_conn_id=SFTP_CONN_ID,
delegate_to=DELEGATE_TO,
)
task.execute(None)
sftp_hook.return_value.get_tree_map.assert_called_with(
"main_dir", prefix="main_dir/test_object", delimiter=".json"
)
sftp_hook.return_value.retrieve_file.assert_has_calls(
[
mock.call("main_dir/test_object3.json", mock.ANY),
mock.call("main_dir/sub_dir/test_object3.json", mock.ANY),
]
)
gcs_hook.return_value.upload.assert_has_calls(
[
mock.call(
bucket_name=TEST_BUCKET,
object_name="destination_dir/test_object3.json",
mime_type=DEFAULT_MIME_TYPE,
filename=mock.ANY,
),
mock.call(
bucket_name=TEST_BUCKET,
object_name="destination_dir/sub_dir/test_object3.json",
mime_type=DEFAULT_MIME_TYPE,
filename=mock.ANY,
),
]
)