fix: catch SSHException to trigger the retry on connection dropped

ToucanToco · Oct 13, 2023 · e874a86 · e874a86
1 parent eac79a3
commit e874a86
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/peakina/io/ftp/ftp_utils.py b/peakina/io/ftp/ftp_utils.py
@@ -189,7 +189,7 @@ def ftp_open(url: str, retry: int = _DEFAULT_MAX_RETRY) -> IO[bytes]:  # type: i
     for i in range(1, retry + 1):
         try:
             return _open(url)
-        except (AttributeError, OSError, ftplib.error_temp) as e:
+        except (AttributeError, OSError, ftplib.error_temp, paramiko.SSHException) as e:
             log = logging.getLogger(__name__)
             # If this occurs, we need to see what's actually inside that dir
             # by listing maxi 15 entries

diff --git a/tests/io/ftp/test_ftp_utils.py b/tests/io/ftp/test_ftp_utils.py
@@ -4,6 +4,7 @@
 import ssl
 from urllib.parse import ParseResult
 
+from paramiko.ssh_exception import SSHException
 from pytest import fixture, raises
 from pytest_mock import MockFixture
 
@@ -68,6 +69,7 @@ def test_retry_open(mocker):
         ftplib.error_temp("421 Could not create socket"),
         AttributeError("'NoneType' object has no attribute 'sendall'"),
         OSError("Random OSError"),
+        SSHException("Random connection dropped error"),
         "ok",
     ]
     mock_sleep = mocker.patch("peakina.io.ftp.ftp_utils.sleep")