From 17eaf8eae30e6e2d4fac1b582495a052b5871415 Mon Sep 17 00:00:00 2001 From: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Date: Wed, 3 Apr 2024 23:04:48 +0300 Subject: [PATCH] Resolve `PT012` in `SFTP` provider tests (#38518) Co-authored-by: Bowrna --- airflow/providers/sftp/triggers/sftp.py | 2 +- pyproject.toml | 6 --- tests/providers/sftp/hooks/test_sftp.py | 13 +++--- tests/providers/sftp/operators/test_sftp.py | 47 +++++++++++---------- tests/providers/sftp/sensors/test_sftp.py | 1 - tests/providers/sftp/triggers/test_sftp.py | 12 ++---- 6 files changed, 35 insertions(+), 46 deletions(-) diff --git a/airflow/providers/sftp/triggers/sftp.py b/airflow/providers/sftp/triggers/sftp.py index 1c117a5a131b0..6a39c95870607 100644 --- a/airflow/providers/sftp/triggers/sftp.py +++ b/airflow/providers/sftp/triggers/sftp.py @@ -134,7 +134,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # Break loop to avoid infinite retries on terminal failure break - yield TriggerEvent({"status": "error", "message": exc}) + yield TriggerEvent({"status": "error", "message": str(exc)}) def _get_async_hook(self) -> SFTPHookAsync: return SFTPHookAsync(sftp_conn_id=self.sftp_conn_id) diff --git a/pyproject.toml b/pyproject.toml index 95f01ea7bb9d6..f3d451336d8ae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -375,12 +375,6 @@ combine-as-imports = true "tests/providers/qdrant/operators/test_qdrant.py" = ["E402"] "tests/providers/snowflake/operators/test_snowflake_sql.py" = ["E402"] -# All the test modules which do not follow PT012 yet -"tests/providers/sftp/hooks/test_sftp.py" = ["PT012"] -"tests/providers/sftp/operators/test_sftp.py" = ["PT012"] -"tests/providers/sftp/sensors/test_sftp.py" = ["PT012"] -"tests/providers/sftp/triggers/test_sftp.py" = ["PT012"] - [tool.ruff.lint.flake8-tidy-imports] # Disallow all relative imports. ban-relative-imports = "all" diff --git a/tests/providers/sftp/hooks/test_sftp.py b/tests/providers/sftp/hooks/test_sftp.py index f914e1da567fe..4f7cec34fb39b 100644 --- a/tests/providers/sftp/hooks/test_sftp.py +++ b/tests/providers/sftp/hooks/test_sftp.py @@ -404,13 +404,12 @@ def test_deprecation_ftp_conn_id(self, mock_get_connection): @mock.patch("airflow.providers.sftp.hooks.sftp.SFTPHook.get_connection") def test_invalid_ssh_hook(self, mock_get_connection): - with pytest.raises(AirflowException, match="ssh_hook must be an instance of SSHHook"): - connection = Connection(conn_id="sftp_default", login="root", host="localhost") - mock_get_connection.return_value = connection - with pytest.warns( - AirflowProviderDeprecationWarning, match=r"Parameter `ssh_hook` is deprecated.*" - ): - SFTPHook(ssh_hook="invalid_hook") + connection = Connection(conn_id="sftp_default", login="root", host="localhost") + mock_get_connection.return_value = connection + with pytest.raises(AirflowException, match="ssh_hook must be an instance of SSHHook"), pytest.warns( + AirflowProviderDeprecationWarning, match=r"Parameter `ssh_hook` is deprecated.*" + ): + SFTPHook(ssh_hook="invalid_hook") @mock.patch("airflow.providers.ssh.hooks.ssh.SSHHook.get_connection") def test_valid_ssh_hook(self, mock_get_connection): diff --git a/tests/providers/sftp/operators/test_sftp.py b/tests/providers/sftp/operators/test_sftp.py index 8f72a41915fbb..984c5e925b987 100644 --- a/tests/providers/sftp/operators/test_sftp.py +++ b/tests/providers/sftp/operators/test_sftp.py @@ -310,14 +310,14 @@ def test_file_transfer_with_intermediate_dir_error_get(self, dag_maker, create_r def test_arg_checking(self): dag = DAG(dag_id="unit_tests_sftp_op_arg_checking", default_args={"start_date": DEFAULT_DATE}) # Exception should be raised if neither ssh_hook nor ssh_conn_id is provided + task_0 = SFTPOperator( + task_id="test_sftp_0", + local_filepath=self.test_local_filepath, + remote_filepath=self.test_remote_filepath, + operation=SFTPOperation.PUT, + dag=dag, + ) with pytest.raises(AirflowException, match="Cannot operate without sftp_hook or ssh_conn_id."): - task_0 = SFTPOperator( - task_id="test_sftp_0", - local_filepath=self.test_local_filepath, - remote_filepath=self.test_remote_filepath, - operation=SFTPOperation.PUT, - dag=dag, - ) task_0.execute(None) # if ssh_hook is invalid/not provided, use ssh_conn_id to create SSHHook @@ -360,31 +360,32 @@ def test_arg_checking(self): task_3.execute(None) assert task_3.sftp_hook.ssh_conn_id == self.hook.ssh_conn_id + task_4 = SFTPOperator( + task_id="test_sftp_4", + local_filepath=self.test_local_filepath, + remote_filepath=self.test_remote_filepath, + operation="invalid_operation", + dag=dag, + ) # Exception should be raised if operation is invalid with pytest.raises(TypeError, match="Unsupported operation value invalid_operation, "): - task_4 = SFTPOperator( - task_id="test_sftp_4", - local_filepath=self.test_local_filepath, - remote_filepath=self.test_remote_filepath, - operation="invalid_operation", - dag=dag, - ) task_4.execute(None) + task_5 = SFTPOperator( + task_id="test_sftp_5", + ssh_hook=self.hook, + sftp_hook=SFTPHook(), + local_filepath=self.test_local_filepath, + remote_filepath=self.test_remote_filepath, + operation=SFTPOperation.PUT, + dag=dag, + ) + # Exception should be raised if both ssh_hook and sftp_hook are provided with pytest.raises( AirflowException, match="Both `ssh_hook` and `sftp_hook` are defined. Please use only one of them.", ): - task_5 = SFTPOperator( - task_id="test_sftp_5", - ssh_hook=self.hook, - sftp_hook=SFTPHook(), - local_filepath=self.test_local_filepath, - remote_filepath=self.test_remote_filepath, - operation=SFTPOperation.PUT, - dag=dag, - ) task_5.execute(None) task_6 = SFTPOperator( diff --git a/tests/providers/sftp/sensors/test_sftp.py b/tests/providers/sftp/sensors/test_sftp.py index fb98b898c8554..6a08b377ce80f 100644 --- a/tests/providers/sftp/sensors/test_sftp.py +++ b/tests/providers/sftp/sensors/test_sftp.py @@ -64,7 +64,6 @@ def test_sftp_failure(self, sftp_hook_mock, soft_fail: bool, expected_exception) context = {"ds": "1970-01-01"} with pytest.raises(expected_exception): sftp_sensor.poke(context) - sftp_hook_mock.return_value.get_mod_time.assert_called_once_with("/path/to/file/1970-01-01.txt") def test_hook_not_created_during_init(self): sftp_sensor = SFTPSensor(task_id="unit_test", path="/path/to/file/1970-01-01.txt") diff --git a/tests/providers/sftp/triggers/test_sftp.py b/tests/providers/sftp/triggers/test_sftp.py index f1bfa52533e2f..5816769129da2 100644 --- a/tests/providers/sftp/triggers/test_sftp.py +++ b/tests/providers/sftp/triggers/test_sftp.py @@ -174,14 +174,10 @@ async def test_sftp_trigger_run_trigger_failure_state(self, mock_get_files_by_pa mock_get_files_by_pattern.side_effect = Exception("An unexpected exception") trigger = SFTPTrigger(path="test/path/", sftp_conn_id="sftp_default", file_pattern="my_test_file") - - expected_event = {"status": "failure", "message": "An unexpected exception"} - - with pytest.raises(Exception): - generator = trigger.run() - actual_event = await generator.asend(None) - - assert TriggerEvent(expected_event) == actual_event + expected_event = {"status": "error", "message": "An unexpected exception"} + generator = trigger.run() + actual_event = await generator.asend(None) + assert TriggerEvent(expected_event) == actual_event @pytest.mark.asyncio @mock.patch("airflow.providers.sftp.hooks.sftp.SFTPHookAsync.get_files_and_attrs_by_pattern")