From a18f8d2ff3581f555403853a229e1937bd70ecf1 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Tue, 11 Oct 2022 23:50:07 +0200
Subject: [PATCH] Add error handling to LaMetric button platform (#80136)

---
 homeassistant/components/lametric/button.py  |  2 +
 homeassistant/components/lametric/helpers.py | 46 ++++++++++++++++
 tests/components/lametric/test_button.py     | 58 +++++++++++++++++++-
 3 files changed, 105 insertions(+), 1 deletion(-)
 create mode 100644 homeassistant/components/lametric/helpers.py

diff --git a/homeassistant/components/lametric/button.py b/homeassistant/components/lametric/button.py
index 4d8c75f0ab0..8de7ddb16ff 100644
--- a/homeassistant/components/lametric/button.py
+++ b/homeassistant/components/lametric/button.py
@@ -16,6 +16,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from .const import DOMAIN
 from .coordinator import LaMetricDataUpdateCoordinator
 from .entity import LaMetricEntity
+from .helpers import lametric_exception_handler
 
 
 @dataclass
@@ -81,6 +82,7 @@ class LaMetricButtonEntity(LaMetricEntity, ButtonEntity):
         self.entity_description = description
         self._attr_unique_id = f"{coordinator.data.serial_number}-{description.key}"
 
+    @lametric_exception_handler
     async def async_press(self) -> None:
         """Send out a command to LaMetric."""
         await self.entity_description.press_fn(self.coordinator.lametric)
diff --git a/homeassistant/components/lametric/helpers.py b/homeassistant/components/lametric/helpers.py
new file mode 100644
index 00000000000..e9b1b941dae
--- /dev/null
+++ b/homeassistant/components/lametric/helpers.py
@@ -0,0 +1,46 @@
+"""Helpers for LaMetric."""
+from __future__ import annotations
+
+from collections.abc import Callable, Coroutine
+from typing import Any, TypeVar
+
+from demetriek import LaMetricConnectionError, LaMetricError
+from typing_extensions import Concatenate, ParamSpec
+
+from homeassistant.exceptions import HomeAssistantError
+
+from .entity import LaMetricEntity
+
+_LaMetricEntityT = TypeVar("_LaMetricEntityT", bound=LaMetricEntity)
+_P = ParamSpec("_P")
+
+
+def lametric_exception_handler(
+    func: Callable[Concatenate[_LaMetricEntityT, _P], Coroutine[Any, Any, Any]]
+) -> Callable[Concatenate[_LaMetricEntityT, _P], Coroutine[Any, Any, None]]:
+    """Decorate LaMetric calls to handle LaMetric exceptions.
+
+    A decorator that wraps the passed in function, catches LaMetric errors,
+    and handles the availability of the device in the data coordinator.
+    """
+
+    async def handler(
+        self: _LaMetricEntityT, *args: _P.args, **kwargs: _P.kwargs
+    ) -> None:
+        try:
+            await func(self, *args, **kwargs)
+            self.coordinator.async_update_listeners()
+
+        except LaMetricConnectionError as error:
+            self.coordinator.last_update_success = False
+            self.coordinator.async_update_listeners()
+            raise HomeAssistantError(
+                "Error communicating with the LaMetric device"
+            ) from error
+
+        except LaMetricError as error:
+            raise HomeAssistantError(
+                "Invalid response from the LaMetric device"
+            ) from error
+
+    return handler
diff --git a/tests/components/lametric/test_button.py b/tests/components/lametric/test_button.py
index cd55c9914f5..d37b6dd1c18 100644
--- a/tests/components/lametric/test_button.py
+++ b/tests/components/lametric/test_button.py
@@ -1,12 +1,19 @@
 """Tests for the LaMetric button platform."""
 from unittest.mock import MagicMock
 
+from demetriek import LaMetricConnectionError, LaMetricError
 import pytest
 
 from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, SERVICE_PRESS
 from homeassistant.components.lametric.const import DOMAIN
-from homeassistant.const import ATTR_ENTITY_ID, ATTR_ICON, STATE_UNKNOWN
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    ATTR_ICON,
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.entity import EntityCategory
 
@@ -111,3 +118,52 @@ async def test_button_app_previous(
     state = hass.states.get("button.frenck_s_lametric_previous_app")
     assert state
     assert state.state == "2022-09-19T12:07:30+00:00"
+
+
+@pytest.mark.freeze_time("2022-10-11 22:00:00")
+async def test_button_error(
+    hass: HomeAssistant,
+    init_integration: MockConfigEntry,
+    mock_lametric: MagicMock,
+) -> None:
+    """Test error handling of the LaMetric buttons."""
+    mock_lametric.app_next.side_effect = LaMetricError
+
+    with pytest.raises(
+        HomeAssistantError, match="Invalid response from the LaMetric device"
+    ):
+        await hass.services.async_call(
+            BUTTON_DOMAIN,
+            SERVICE_PRESS,
+            {ATTR_ENTITY_ID: "button.frenck_s_lametric_next_app"},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+
+    state = hass.states.get("button.frenck_s_lametric_next_app")
+    assert state
+    assert state.state == "2022-10-11T22:00:00+00:00"
+
+
+async def test_button_connection_error(
+    hass: HomeAssistant,
+    init_integration: MockConfigEntry,
+    mock_lametric: MagicMock,
+) -> None:
+    """Test connection error handling of the LaMetric buttons."""
+    mock_lametric.app_next.side_effect = LaMetricConnectionError
+
+    with pytest.raises(
+        HomeAssistantError, match="Error communicating with the LaMetric device"
+    ):
+        await hass.services.async_call(
+            BUTTON_DOMAIN,
+            SERVICE_PRESS,
+            {ATTR_ENTITY_ID: "button.frenck_s_lametric_next_app"},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+
+    state = hass.states.get("button.frenck_s_lametric_next_app")
+    assert state
+    assert state.state == STATE_UNAVAILABLE