def test_execute_copy_single_file(self, sftp_hook, gcs_hook): task = GCSToSFTPOperator( task_id=TASK_ID, source_bucket=TEST_BUCKET, source_object=SOURCE_OBJECT_NO_WILDCARD, destination_path=DESTINATION_SFTP, move_object=False, gcp_conn_id=GCP_CONN_ID, sftp_conn_id=SFTP_CONN_ID, delegate_to=DELEGATE_TO, impersonation_chain=IMPERSONATION_CHAIN, ) task.execute({}) 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) args, kwargs = gcs_hook.return_value.download.call_args assert kwargs["bucket_name"] == TEST_BUCKET assert kwargs["object_name"] == SOURCE_OBJECT_NO_WILDCARD args, kwargs = sftp_hook.return_value.store_file.call_args assert args[0] == os.path.join(DESTINATION_SFTP, SOURCE_OBJECT_NO_WILDCARD) gcs_hook.return_value.delete.assert_not_called()
def test_execute_move_single_file(self, sftp_hook, gcs_hook): task = GCSToSFTPOperator( 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_move_single_file( self, source_object, target_object, keep_directory_structure, sftp_hook_mock, gcs_hook_mock ): task = GCSToSFTPOperator( task_id=TASK_ID, source_bucket=TEST_BUCKET, source_object=source_object, destination_path=DESTINATION_SFTP, keep_directory_structure=keep_directory_structure, 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_mock.assert_called_once_with( gcp_conn_id=GCP_CONN_ID, delegate_to=DELEGATE_TO, impersonation_chain=IMPERSONATION_CHAIN, ) sftp_hook_mock.assert_called_once_with(SFTP_CONN_ID) gcs_hook_mock.return_value.download.assert_called_with( bucket_name=TEST_BUCKET, object_name=source_object, filename=mock.ANY ) sftp_hook_mock.return_value.store_file.assert_called_with( os.path.join(DESTINATION_SFTP, target_object), mock.ANY ) gcs_hook_mock.return_value.delete.assert_called_once_with(TEST_BUCKET, source_object)
def test_execute_more_than_one_wildcard_exception(self, sftp_hook_mock, gcs_hook_mock): operator = GCSToSFTPOperator( task_id=TASK_ID, source_bucket=TEST_BUCKET, source_object="csv/*/test_*.csv", destination_path=DESTINATION_SFTP, move_object=False, gcp_conn_id=GCP_CONN_ID, sftp_conn_id=SFTP_CONN_ID, delegate_to=DELEGATE_TO, ) with pytest.raises(AirflowException): operator.execute(None)
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 = GCSToSFTPOperator( 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 pytest.raises(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 = GCSToSFTPOperator( 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 assert call_one[0] == (TEST_BUCKET, "test_object/file1.txt") assert call_two[0] == (TEST_BUCKET, "test_object/file2.txt")
def test_execute_move_with_wildcard( self, source_object, prefix, delimiter, gcs_files_list, target_objects, keep_directory_structure, sftp_hook_mock, gcs_hook_mock, ): gcs_hook_mock.return_value.list.return_value = gcs_files_list operator = GCSToSFTPOperator( task_id=TASK_ID, source_bucket=TEST_BUCKET, source_object=source_object, destination_path=DESTINATION_SFTP, keep_directory_structure=keep_directory_structure, move_object=True, gcp_conn_id=GCP_CONN_ID, sftp_conn_id=SFTP_CONN_ID, delegate_to=DELEGATE_TO, ) operator.execute(None) gcs_hook_mock.return_value.list.assert_called_with(TEST_BUCKET, delimiter=delimiter, prefix=prefix) gcs_hook_mock.return_value.download.assert_has_calls( [ mock.call(bucket_name=TEST_BUCKET, object_name=gcs_file, filename=mock.ANY) for gcs_file in gcs_files_list ] ) sftp_hook_mock.return_value.store_file.assert_has_calls( [ mock.call(os.path.join(DESTINATION_SFTP, target_object), mock.ANY) for target_object in target_objects ] ) gcs_hook_mock.return_value.delete.assert_has_calls( [mock.call(TEST_BUCKET, gcs_file) for gcs_file in gcs_files_list] )