Retry sftp_storage setup when the SSH connection fails (#181769)

homeassistant/components/sftp_storage/__init__.py

@@ -10,9 +10,9 @@ from homeassistant.components.backup import BackupAgentError
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryError
+from homeassistant.exceptions import ConfigEntryError, ConfigEntryNotReady
 
-from .client import BackupAgentClient
+from .client import BackupAgentClient, SFTPConnectionError
 from .const import (
     CONF_BACKUP_LOCATION,
     CONF_PRIVATE_KEY_FILE,
@@ -55,6 +55,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: SFTPConfigEntry) -> bool
     try:
         client = BackupAgentClient(entry, hass)
         await client.open()
+    except SFTPConnectionError as e:
+        raise ConfigEntryNotReady(str(e)) from e
     except BackupAgentError as e:
         raise ConfigEntryError from e
 

homeassistant/components/sftp_storage/client.py

@@ -13,7 +13,7 @@ from asyncssh import (
     SSHClientConnectionOptions,
     connect,
 )
-from asyncssh.misc import PermissionDenied
+from asyncssh.misc import Error as SSHError, PermissionDenied
 from asyncssh.sftp import SFTPNoSuchFile, SFTPPermissionDenied
 
 from homeassistant.components.backup import (
@@ -29,6 +29,10 @@ if TYPE_CHECKING:
     from . import SFTPConfigEntry, SFTPConfigEntryData
 
 
+class SFTPConnectionError(BackupAgentError):
+    """Error raised when the SSH connection could not be established."""
+
+
 def get_client_options(cfg: SFTPConfigEntryData) -> SSHClientConnectionOptions:
     """Get `SSHClientConnectionOptions` for use with `hass.async_add_executor_job`."""
 
@@ -295,12 +299,17 @@ class BackupAgentClient:
                     get_client_options, self.cfg.runtime_data
                 ),
             )
-        except (OSError, PermissionDenied) as e:
+        except PermissionDenied as e:
             raise BackupAgentError(
                 "Failure while attempting to establish SSH"
                 " connection. Please check SSH credentials"
                 " and if changed, re-install the integration"
             ) from e
+        except (OSError, SSHError) as e:
+            raise SFTPConnectionError(
+                f"Failed to establish SSH connection to"
+                f" {self.cfg.runtime_data.host}: {e}"
+            ) from e
 
         # Configure SFTP Client Connection
         try:
@@ -311,5 +320,9 @@ class BackupAgentClient:
                 "Failed to create SFTP client."
                 " Re-installing integration might be required"
             ) from e
+        except (OSError, SSHError) as e:
+            raise SFTPConnectionError(
+                f"Failed to open SFTP session on {self.cfg.runtime_data.host}: {e}"
+            ) from e
 
         return self

tests/components/sftp_storage/test_init.py

@@ -3,7 +3,8 @@
 from pathlib import Path
 from unittest.mock import patch
 
-from asyncssh.sftp import SFTPPermissionDenied
+from asyncssh.misc import ChannelOpenError, ConnectionLost, PermissionDenied
+from asyncssh.sftp import SFTPConnectionLost, SFTPPermissionDenied
 import pytest
 
 from homeassistant.components.sftp_storage import SFTPConfigEntryData
@@ -73,15 +74,61 @@ async def test_setup_error(
     assert entries[0].state is ConfigEntryState.SETUP_ERROR
 
 
-async def test_setup_unexpected_error(
+@pytest.mark.parametrize(
+    "connect_error",
+    [OSError("Error message"), ConnectionLost("Connection lost")],
+    ids=["oserror", "connection_lost"],
+)
+async def test_setup_connection_error_is_retried(
+    hass: HomeAssistant,
+    setup_integration: ComponentSetup,
+    caplog: pytest.LogCaptureFixture,
+    connect_error: Exception,
+) -> None:
+    """Test that a connection failure leaves the entry in a retrying state."""
+    with patch(
+        "homeassistant.components.sftp_storage.client.connect",
+        side_effect=connect_error,
+    ):
+        await setup_integration()
+
+    entries = hass.config_entries.async_entries(DOMAIN)
+    assert len(entries) == 1
+    assert entries[0].state is ConfigEntryState.SETUP_RETRY
+    assert "Failed to establish SSH connection to" in caplog.text
+
+
+@pytest.mark.parametrize(
+    "sftp_error",
+    [ChannelOpenError(1, "Channel open failed"), SFTPConnectionLost("Connection lost")],
+    ids=["channel_open_error", "sftp_connection_lost"],
+)
+async def test_setup_sftp_session_error_is_retried(
+    mock_ssh_connection: SSHClientConnectionMock,
+    hass: HomeAssistant,
+    setup_integration: ComponentSetup,
+    caplog: pytest.LogCaptureFixture,
+    sftp_error: Exception,
+) -> None:
+    """Test that losing the session after connecting is also retried."""
+    mock_ssh_connection._sftp._mock_chdir.side_effect = sftp_error
+    await setup_integration()
+
+    entries = hass.config_entries.async_entries(DOMAIN)
+    assert len(entries) == 1
+    assert entries[0].state is ConfigEntryState.SETUP_RETRY
+    assert "Failed to open SFTP session on" in caplog.text
+
+
+async def test_setup_invalid_credentials(
     hass: HomeAssistant,
     setup_integration: ComponentSetup,
     caplog: pytest.LogCaptureFixture,
 ) -> None:
-    """Test setup error."""
+    """Test that rejected credentials are not retried."""
     with patch(
         "homeassistant.components.sftp_storage.client.connect",
-        side_effect=OSError("Error message"),
+        side_effect=PermissionDenied("Permission denied"),
     ):
         await setup_integration()