RachitSharma2001 commented on code in PR #28318:
URL: https://github.com/apache/airflow/pull/28318#discussion_r1047889098
##########
tests/providers/ftp/operators/test_ftp.py:
##########
@@ -184,3 +184,112 @@ def test_unequal_local_remote_file_paths(self):
local_filepath=["/tmp/test1", "/tmp/test2"],
remote_filepath="/tmp/test1",
)
+
+
+class TestFTPSFileTransmitOperator:
+ def setup_method(self):
+ self.test_local_dir = "/tmp"
+ self.test_local_dir_int = "/tmp/interdir"
+ self.test_remote_dir = "/ftpshome"
+ self.test_remote_dir_int = "/ftpshome/interdir"
+ self.test_local_filename = "test_local_file"
+ self.test_remote_filename = "test_remote_file"
+ self.test_local_filepath =
f"{self.test_local_dir}/{self.test_local_filename}"
+ self.test_remote_filepath =
f"{self.test_remote_dir}/{self.test_remote_filename}"
+ self.test_local_filepath_int_dir =
f"{self.test_local_dir_int}/{self.test_local_filename}"
+ self.test_remote_filepath_int_dir =
f"{self.test_remote_dir_int}/{self.test_remote_filename}"
+
+ def teardown_method(self):
+ if os.path.exists(self.test_local_dir_int):
+ os.rmdir(self.test_local_dir_int)
+
+ @mock.patch("airflow.providers.ftp.operators.ftp.FTPSHook.store_file")
+
@mock.patch("airflow.providers.ftp.operators.ftp.FTPSHook.create_directory")
+ def test_file_transfer_put(self, mock_create_dir, mock_put):
+ ftps_op = FTPSFileTransmitOperator(
+ task_id="test_ftps_put",
+ ftp_conn_id=DEFAULT_CONN_ID,
+ local_filepath=self.test_local_filepath,
+ remote_filepath=self.test_remote_filepath,
+ operation=FTPOperation.PUT,
+ )
+ ftps_op.execute(None)
+ assert not mock_create_dir.called
+ mock_put.assert_called_once_with(self.test_remote_filepath,
self.test_local_filepath)
+
+ @mock.patch("airflow.providers.ftp.operators.ftp.FTPSHook.store_file")
+
@mock.patch("airflow.providers.ftp.operators.ftp.FTPSHook.create_directory")
+ def test_file_transfer_with_intermediate_dir_put(self, mock_create_dir,
mock_put):
+ ftps_op = FTPSFileTransmitOperator(
+ task_id="test_ftps_put_imm_dirs",
+ ftp_conn_id=DEFAULT_CONN_ID,
+ local_filepath=self.test_local_filepath,
+ remote_filepath=self.test_remote_filepath_int_dir,
+ operation=FTPOperation.PUT,
+ create_intermediate_dirs=True,
+ )
+ ftps_op.execute(None)
+ mock_create_dir.assert_called_with(self.test_remote_dir_int)
+ mock_put.assert_called_once_with(self.test_remote_filepath_int_dir,
self.test_local_filepath)
+
+ @mock.patch("airflow.providers.ftp.operators.ftp.FTPSHook.retrieve_file")
+ def test_file_transfer_get(self, mock_get):
+ ftps_op = FTPSFileTransmitOperator(
+ task_id="test_ftps_get",
+ ftp_conn_id=DEFAULT_CONN_ID,
+ local_filepath=self.test_local_filepath,
+ remote_filepath=self.test_remote_filepath,
+ operation=FTPOperation.GET,
+ )
+ ftps_op.execute(None)
+ mock_get.assert_called_once_with(self.test_remote_filepath,
self.test_local_filepath)
+
+ @mock.patch("airflow.providers.ftp.operators.ftp.FTPSHook.retrieve_file")
+ def test_file_transfer_with_intermediate_dir_get(self, mock_get):
+ ftps_op = FTPSFileTransmitOperator(
+ task_id="test_ftps_get_imm_dirs",
+ ftp_conn_id=DEFAULT_CONN_ID,
+ local_filepath=self.test_local_filepath_int_dir,
+ remote_filepath=self.test_remote_filepath,
+ operation=FTPOperation.GET,
+ create_intermediate_dirs=True,
+ )
+ ftps_op.execute(None)
+ assert os.path.exists(self.test_local_dir_int)
+ mock_get.assert_called_once_with(self.test_remote_filepath,
self.test_local_filepath_int_dir)
+
+ @mock.patch("airflow.providers.ftp.operators.ftp.FTPSHook.retrieve_file")
+ def test_multiple_paths_get(self, mock_get):
+ local_filepath = ["/tmp/ltest1", "/tmp/ltest2"]
+ remote_filepath = ["/tmp/rtest1", "/tmp/rtest2"]
+ ftps_op = FTPSFileTransmitOperator(
+ task_id="test_multiple_paths_get",
+ ftp_conn_id=DEFAULT_CONN_ID,
+ local_filepath=local_filepath,
+ remote_filepath=remote_filepath,
+ operation=FTPOperation.GET,
+ )
+ ftps_op.execute(None)
+ assert mock_get.call_count == 2
+ args0, _ = mock_get.call_args_list[0]
+ args1, _ = mock_get.call_args_list[1]
+ assert args0 == (remote_filepath[0], local_filepath[0])
+ assert args1 == (remote_filepath[1], local_filepath[1])
Review Comment:
@ferruzzi I've added this change as well. Let me know if it looks good and
if any other change is needed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
