test(sftp): fix broken tests

apache · Apr 10, 2024 · 119fd3b · 119fd3b
1 parent 4cfbe96
commit 119fd3b
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/tests/providers/sftp/hooks/test_sftp.py b/tests/providers/sftp/hooks/test_sftp.py
@@ -712,52 +712,57 @@ async def test_list_directory_path_does_not_exist(self, mock_hook_get_conn):
         """
         Assert that AirflowException is raised when path does not exist on SFTP server
         """
-        mock_hook_get_conn.return_value = MockSSHClient()
+        mock_hook_get_conn.return_value.__aenter__.return_value = MockSSHClient()
+
         hook = SFTPHookAsync()
 
         expected_files = None
         files = await hook.list_directory(path="/path/does_not/exist/")
         assert files == expected_files
+        mock_hook_get_conn.return_value.__aexit__.assert_called()
 
     @patch("airflow.providers.sftp.hooks.sftp.SFTPHookAsync._get_conn")
     @pytest.mark.asyncio
     async def test_read_directory_path_does_not_exist(self, mock_hook_get_conn):
         """
         Assert that AirflowException is raised when path does not exist on SFTP server
         """
-        mock_hook_get_conn.return_value = MockSSHClient()
+        mock_hook_get_conn.return_value.__aenter__.return_value = MockSSHClient()
         hook = SFTPHookAsync()
 
         expected_files = None
         files = await hook.read_directory(path="/path/does_not/exist/")
         assert files == expected_files
+        mock_hook_get_conn.return_value.__aexit__.assert_called()
 
     @patch("airflow.providers.sftp.hooks.sftp.SFTPHookAsync._get_conn")
     @pytest.mark.asyncio
     async def test_list_directory_path_has_files(self, mock_hook_get_conn):
         """
         Assert that file list is returned when path exists on SFTP server
         """
-        mock_hook_get_conn.return_value = MockSSHClient()
+        mock_hook_get_conn.return_value.__aenter__.return_value = MockSSHClient()
         hook = SFTPHookAsync()
 
         expected_files = ["..", ".", "file"]
         files = await hook.list_directory(path="/path/exists/")
         assert sorted(files) == sorted(expected_files)
+        mock_hook_get_conn.return_value.__aexit__.assert_called()
 
     @patch("airflow.providers.sftp.hooks.sftp.SFTPHookAsync._get_conn")
     @pytest.mark.asyncio
     async def test_get_file_by_pattern_with_match(self, mock_hook_get_conn):
         """
         Assert that filename is returned when file pattern is matched on SFTP server
         """
-        mock_hook_get_conn.return_value = MockSSHClient()
+        mock_hook_get_conn.return_value.__aenter__.return_value = MockSSHClient()
         hook = SFTPHookAsync()
 
         files = await hook.get_files_and_attrs_by_pattern(path="/path/exists/", fnmatch_pattern="file")
 
         assert len(files) == 1
         assert files[0].filename == "file"
+        mock_hook_get_conn.return_value.__aexit__.assert_called()
 
     @pytest.mark.asyncio
     @patch("airflow.providers.sftp.hooks.sftp.SFTPHookAsync._get_conn")