Fix resp may be undefined end_backup issue (#5224)

home-assistant · Aug 5, 2024 · 17ee234 · 17ee234
1 parent 61034df
commit 17ee234
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 13 deletions.
diff --git a/supervisor/homeassistant/module.py b/supervisor/homeassistant/module.py
@@ -373,12 +373,12 @@ async def end_backup(self) -> None:
             _LOGGER.warning(
                 "Error resuming normal operations after backup of Home Assistant Core. Check HA Core logs."
             )
-
-        if resp and not resp.get(ATTR_SUCCESS):
-            _LOGGER.warning(
-                "Error resuming normal operations after backup of Home Assistant Core due to: %s. Check HA Core logs.",
-                resp.get(ATTR_ERROR, {}).get(ATTR_MESSAGE, ""),
-            )
+        else:
+            if resp and not resp.get(ATTR_SUCCESS):
+                _LOGGER.warning(
+                    "Error resuming normal operations after backup of Home Assistant Core due to: %s. Check HA Core logs.",
+                    resp.get(ATTR_ERROR, {}).get(ATTR_MESSAGE, ""),
+                )
 
     @Job(name="home_assistant_module_backup")
     async def backup(

diff --git a/tests/homeassistant/test_module.py b/tests/homeassistant/test_module.py
@@ -5,12 +5,17 @@
 from pathlib import Path
 from unittest.mock import AsyncMock, patch
 
-from pytest import LogCaptureFixture
+from pytest import LogCaptureFixture, raises
 
 from supervisor.const import CoreState
 from supervisor.coresys import CoreSys
 from supervisor.docker.interface import DockerInterface
+from supervisor.exceptions import (
+    HomeAssistantBackupError,
+    HomeAssistantWSConnectionError,
+)
 from supervisor.homeassistant.secrets import HomeAssistantSecrets
+from supervisor.homeassistant.websocket import HomeAssistantWebSocket
 
 
 async def test_load(
@@ -21,12 +26,14 @@ async def test_load(
         secrets.write("hello: world\n")
 
     # Unwrap read_secrets to prevent throttling between tests
-    with patch.object(DockerInterface, "attach") as attach, patch.object(
-        DockerInterface, "check_image"
-    ) as check_image, patch.object(
-        HomeAssistantSecrets,
-        "_read_secrets",
-        new=HomeAssistantSecrets._read_secrets.__wrapped__,
+    with (
+        patch.object(DockerInterface, "attach") as attach,
+        patch.object(DockerInterface, "check_image") as check_image,
+        patch.object(
+            HomeAssistantSecrets,
+            "_read_secrets",
+            new=HomeAssistantSecrets._read_secrets.__wrapped__,
+        ),
     ):
         await coresys.homeassistant.load()
 
@@ -70,3 +77,34 @@ def test_write_pulse_error(coresys: CoreSys, caplog: LogCaptureFixture):
 
         assert "can't write pulse/client.config" in caplog.text
         assert coresys.core.healthy is False
+
+
+async def test_begin_backup_ws_error(coresys: CoreSys):
+    """Test WS error when beginning backup."""
+    # pylint: disable-next=protected-access
+    coresys.homeassistant.websocket._client.async_send_command.side_effect = (
+        HomeAssistantWSConnectionError
+    )
+    with (
+        patch.object(HomeAssistantWebSocket, "_can_send", return_value=True),
+        raises(
+            HomeAssistantBackupError,
+            match="Preparing backup of Home Assistant Core failed. Check HA Core logs.",
+        ),
+    ):
+        await coresys.homeassistant.begin_backup()
+
+
+async def test_end_backup_ws_error(coresys: CoreSys, caplog: LogCaptureFixture):
+    """Test WS error when ending backup."""
+    # pylint: disable-next=protected-access
+    coresys.homeassistant.websocket._client.async_send_command.side_effect = (
+        HomeAssistantWSConnectionError
+    )
+    with patch.object(HomeAssistantWebSocket, "_can_send", return_value=True):
+        await coresys.homeassistant.end_backup()
+
+    assert (
+        "Error resuming normal operations after backup of Home Assistant Core. Check HA Core logs."
+        in caplog.text
+    )