Improve exceptions handling in WLED

homeassistant/components/wled/config_flow.py

@@ -3,7 +3,14 @@
 from typing import Any
 
 import voluptuous as vol
-from wled import WLED, Device, WLEDConnectionError, WLEDUnsupportedVersionError
+from wled import (
+    WLED,
+    Device,
+    WLEDConnectionError,
+    WLEDEmptyResponseError,
+    WLEDInvalidResponseError,
+    WLEDUnsupportedVersionError,
+)
 import yarl
 
 from homeassistant.components import onboarding
@@ -62,6 +69,11 @@ async def async_step_user(
                 errors["base"] = "unsupported_version"
             except WLEDConnectionError:
                 errors["base"] = "cannot_connect"
+            except (WLEDInvalidResponseError, WLEDEmptyResponseError) as er:
+                if "presets" in str(er):
+                    errors["base"] = "invalid_response_presets"
+                else:
+                    errors["base"] = "invalid_response"
             else:
                 mac_address = normalize_mac_address(device.info.mac_address)
                 await self.async_set_unique_id(mac_address, raise_on_progress=False)
@@ -119,6 +131,14 @@ async def async_step_zeroconf(
             self.discovered_device = await self._async_get_device(discovery_info.host)
         except WLEDUnsupportedVersionError:
             return self.async_abort(reason="unsupported_version")
+        except (WLEDInvalidResponseError, WLEDEmptyResponseError) as ex:
+            return self.async_abort(
+                reason=(
+                    "invalid_response_presets"
+                    if "presets" in str(ex)
+                    else "invalid_response"
+                ),
+            )
         except WLEDConnectionError:
             return self.async_abort(reason="cannot_connect")
 

homeassistant/components/wled/coordinator.py

@@ -7,7 +7,9 @@
     Device as WLEDDevice,
     Releases,
     WLEDConnectionClosedError,
+    WLEDEmptyResponseError,
     WLEDError,
+    WLEDInvalidResponseError,
     WLEDReleases,
     WLEDUnsupportedVersionError,
 )
@@ -154,13 +156,23 @@ async def _async_update_data(self) -> WLEDDevice:
                 translation_key="unsupported_version",
                 translation_placeholders={"error": str(error)},
             ) from error
+        except (WLEDInvalidResponseError, WLEDEmptyResponseError) as error:
+            translation_key = (
+                "invalid_response_presets_wled_error"
+                if "presets" in str(error)
+                else "invalid_response_wled_error"
+            )
+            raise UpdateFailed(
+                translation_domain=DOMAIN,
+                translation_key=translation_key,
+                translation_placeholders={"error": str(error)},
+            ) from error
         except WLEDError as error:
             raise UpdateFailed(
                 translation_domain=DOMAIN,
                 translation_key="invalid_response_wled_error",
                 translation_placeholders={"error": str(error)},
             ) from error
-
         device_mac_address = normalize_mac_address(device.info.mac_address)
         if device_mac_address != self.config_mac_address:
             raise ConfigEntryError(

homeassistant/components/wled/helpers.py

@@ -3,7 +3,12 @@
 from collections.abc import Callable, Coroutine
 from typing import Any, Concatenate
 
-from wled import WLEDConnectionError, WLEDError
+from wled import (
+    WLEDConnectionError,
+    WLEDEmptyResponseError,
+    WLEDError,
+    WLEDInvalidResponseError,
+)
 
 from homeassistant.exceptions import HomeAssistantError
 
@@ -33,6 +38,17 @@ async def handler(self: _WLEDEntityT, *args: _P.args, **kwargs: _P.kwargs) -> No
                 translation_key="connection_error",
                 translation_placeholders={"error": str(error)},
             ) from error
+        except (WLEDInvalidResponseError, WLEDEmptyResponseError) as error:
+            translation_key = (
+                "invalid_response_presets_wled_error"
+                if "presets" in str(error)
+                else "invalid_response_wled_error"
+            )
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key=translation_key,
+                translation_placeholders={"error": str(error)},
+            ) from error
         except WLEDError as error:
             raise HomeAssistantError(
                 translation_domain=DOMAIN,

homeassistant/components/wled/strings.json

@@ -12,6 +12,8 @@
     },
     "error": {
       "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
+      "invalid_response": "Received an unexpected response from the device. Please check your device and try again.",
+      "invalid_response_presets": "Failed to download presets from the device. Check preset configurations in WLED UI.",
       "unsupported_version": "[%key:component::wled::common::unsupported_version%]"
     },
     "flow_title": "{name}",
@@ -139,9 +141,15 @@
     "connection_error": {
       "message": "Error communicating with WLED API: {error}"
     },
+    "install_update_wled_error": {
+      "message": "Error installing WLED update: {error}"
+    },
     "invalid_response_github_error": {
       "message": "Invalid response from GitHub API: {error}"
     },
+    "invalid_response_presets_wled_error": {
+      "message": "Failed to download presets from the device. Check preset configurations in WLED UI. Error details: {error}"
+    },
     "invalid_response_wled_error": {
       "message": "Invalid response from WLED API: {error}"
     },

homeassistant/components/wled/update.py

@@ -2,15 +2,19 @@
 
 from typing import Any, cast
 
+from wled import WLEDUpgradeError
+
 from homeassistant.components.update import (
     UpdateDeviceClass,
     UpdateEntity,
     UpdateEntityFeature,
 )
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
 from . import WLED_KEY
+from .const import DOMAIN
 from .coordinator import (
     WLEDConfigEntry,
     WLEDDataUpdateCoordinator,
@@ -110,5 +114,12 @@ async def async_install(
         if version is None:
             # We cast here, as we know that the latest_version is a string.
             version = cast(str, self.latest_version)
-        await self.coordinator.wled.upgrade(version=version)
-        await self.coordinator.async_refresh()
+        try:
+            await self.coordinator.wled.upgrade(version=version)
+            await self.coordinator.async_refresh()
+        except WLEDUpgradeError as error:
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="install_update_wled_error",
+                translation_placeholders={"error": str(error)},
+            ) from error

tests/components/wled/test_button.py

@@ -5,7 +5,12 @@
 
 import pytest
 from syrupy.assertion import SnapshotAssertion
-from wled import WLEDConnectionError, WLEDError
+from wled import (
+    WLEDConnectionError,
+    WLEDEmptyResponseError,
+    WLEDError,
+    WLEDInvalidResponseError,
+)
 
 from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, SERVICE_PRESS
 from homeassistant.components.wled.const import DOMAIN
@@ -87,6 +92,34 @@ async def test_button_restart(
     ("side_effect", "expected_state", "expected_translation_key"),
     [
         (WLEDError, "2021-11-04T16:37:00+00:00", "invalid_response_wled_error"),
+        (
+            WLEDInvalidResponseError(
+                "Received a non-UTF-8 response from request: GET /json"
+            ),
+            "2021-11-04T16:37:00+00:00",
+            "invalid_response_wled_error",
+        ),
+        (
+            WLEDInvalidResponseError(
+                "Received a non-UTF-8 response from request: GET /presets.json"
+            ),
+            "2021-11-04T16:37:00+00:00",
+            "invalid_response_presets_wled_error",
+        ),
+        (
+            WLEDEmptyResponseError(
+                "WLED device at X returned an empty API response on full update"
+            ),
+            "2021-11-04T16:37:00+00:00",
+            "invalid_response_wled_error",
+        ),
+        (
+            WLEDEmptyResponseError(
+                "WLED device at X returned an empty API response on presets update"
+            ),
+            "2021-11-04T16:37:00+00:00",
+            "invalid_response_presets_wled_error",
+        ),
         (WLEDConnectionError, STATE_UNAVAILABLE, "connection_error"),
     ],
 )
@@ -116,3 +149,5 @@ async def test_button_restart_errors(
     # Ensure this made the entity unavailable
     assert (state := hass.states.get("button.wled_rgb_light_restart"))
     assert state.state == expected_state
+
+    mock_wled.reset.assert_called_with()

tests/components/wled/test_config_flow.py

@@ -4,7 +4,12 @@
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
-from wled import WLEDConnectionError, WLEDUnsupportedVersionError
+from wled import (
+    WLEDConnectionError,
+    WLEDEmptyResponseError,
+    WLEDInvalidResponseError,
+    WLEDUnsupportedVersionError,
+)
 
 from homeassistant.components.wled.const import CONF_KEEP_MAIN_LIGHT, DOMAIN
 from homeassistant.config_entries import SOURCE_USER, SOURCE_ZEROCONF
@@ -219,6 +224,30 @@ async def test_zeroconf_during_onboarding(
     [
         (WLEDConnectionError, {"base": "cannot_connect"}),
         (WLEDUnsupportedVersionError, {"base": "unsupported_version"}),
+        (
+            WLEDInvalidResponseError(
+                "Received a non-UTF-8 response from request: GET /json"
+            ),
+            {"base": "invalid_response"},
+        ),
+        (
+            WLEDInvalidResponseError(
+                "Received a non-UTF-8 response from request: GET /presets.json"
+            ),
+            {"base": "invalid_response_presets"},
+        ),
+        (
+            WLEDEmptyResponseError(
+                "WLED device at X returned an empty API response on full update"
+            ),
+            {"base": "invalid_response"},
+        ),
+        (
+            WLEDEmptyResponseError(
+                "WLED device at X returned an empty API response on presets update"
+            ),
+            {"base": "invalid_response_presets"},
+        ),
     ],
 )
 async def test_form_submission_errors(
@@ -237,35 +266,45 @@ async def test_form_submission_errors(
     assert result.get("errors") == errors
 
 
-async def test_zeroconf_connection_error(
-    hass: HomeAssistant, mock_wled: MagicMock
-) -> None:
-    """Test we abort zeroconf flow on WLED connection error."""
-    mock_wled.update.side_effect = WLEDConnectionError
-
-    result = await hass.config_entries.flow.async_init(
-        DOMAIN,
-        context={"source": SOURCE_ZEROCONF},
-        data=ZeroconfServiceInfo(
-            ip_address=ip_address("192.168.1.123"),
-            ip_addresses=[ip_address("192.168.1.123")],
-            hostname="example.local.",
-            name="mock_name",
-            port=None,
-            properties={CONF_MAC: "aabbccddeeff"},
-            type="mock_type",
+@pytest.mark.parametrize(
+    ("side_effect", "expected_reason"),
+    [
+        (WLEDConnectionError, "cannot_connect"),
+        (WLEDUnsupportedVersionError, "unsupported_version"),
+        (
+            WLEDInvalidResponseError(
+                "Received a non-UTF-8 response from request: GET /json"
+            ),
+            "invalid_response",
         ),
-    )
-
-    assert result.get("type") is FlowResultType.ABORT
-    assert result.get("reason") == "cannot_connect"
-
-
-async def test_zeroconf_unsupported_version_error(
-    hass: HomeAssistant, mock_wled: MagicMock
+        (
+            WLEDInvalidResponseError(
+                "Received a non-UTF-8 response from request: GET /presets.json"
+            ),
+            "invalid_response_presets",
+        ),
+        (
+            WLEDEmptyResponseError(
+                "WLED device at X returned an empty API response on full update"
+            ),
+            "invalid_response",
+        ),
+        (
+            WLEDEmptyResponseError(
+                "WLED device at X returned an empty API response on presets update"
+            ),
+            "invalid_response_presets",
+        ),
+    ],
+)
+async def test_zeroconf_errors(
+    hass: HomeAssistant,
+    mock_wled: MagicMock,
+    side_effect: Exception,
+    expected_reason: str,
 ) -> None:
-    """Test we abort zeroconf flow on WLED unsupported version error."""
-    mock_wled.update.side_effect = WLEDUnsupportedVersionError
+    """Test we abort zeroconf flow on WLED errors."""
+    mock_wled.update.side_effect = side_effect
 
     result = await hass.config_entries.flow.async_init(
         DOMAIN,
@@ -282,7 +321,7 @@ async def test_zeroconf_unsupported_version_error(
     )
 
     assert result.get("type") is FlowResultType.ABORT
-    assert result.get("reason") == "unsupported_version"
+    assert result.get("reason") == expected_reason
 
 
 @pytest.mark.usefixtures("mock_wled")