Change update after button press for lamarzocco (#129616)

homeassistant/components/lamarzocco/button.py

@@ -1,11 +1,11 @@
 """Button platform for La Marzocco espresso machines."""
 
+import asyncio
 from collections.abc import Callable, Coroutine
 from dataclasses import dataclass
 from typing import Any
 
 from pylamarzocco.exceptions import RequestNotSuccessful
-from pylamarzocco.lm_machine import LaMarzoccoMachine
 
 from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
 from homeassistant.core import HomeAssistant
@@ -13,9 +13,11 @@ from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .coordinator import LaMarzoccoConfigEntry
+from .coordinator import LaMarzoccoConfigEntry, LaMarzoccoUpdateCoordinator
 from .entity import LaMarzoccoEntity, LaMarzoccoEntityDescription
 
+BACKFLUSH_ENABLED_DURATION = 15
+
 
 @dataclass(frozen=True, kw_only=True)
 class LaMarzoccoButtonEntityDescription(
@@ -24,14 +26,25 @@ class LaMarzoccoButtonEntityDescription(
 ):
     """Description of a La Marzocco button."""
 
-    press_fn: Callable[[LaMarzoccoMachine], Coroutine[Any, Any, None]]
+    press_fn: Callable[[LaMarzoccoUpdateCoordinator], Coroutine[Any, Any, None]]
+
+
+async def async_backflush_and_update(coordinator: LaMarzoccoUpdateCoordinator) -> None:
+    """Press backflush button."""
+    await coordinator.device.start_backflush()
+    # lib will set state optimistically
+    coordinator.async_set_updated_data(None)
+    # backflush is enabled for 15 seconds
+    # then turns off automatically
+    await asyncio.sleep(BACKFLUSH_ENABLED_DURATION + 1)
+    await coordinator.async_request_refresh()
 
 
 ENTITIES: tuple[LaMarzoccoButtonEntityDescription, ...] = (
     LaMarzoccoButtonEntityDescription(
         key="start_backflush",
         translation_key="start_backflush",
-        press_fn=lambda machine: machine.start_backflush(),
+        press_fn=async_backflush_and_update,
     ),
 )
 
@@ -59,7 +72,7 @@ class LaMarzoccoButtonEntity(LaMarzoccoEntity, ButtonEntity):
     async def async_press(self) -> None:
         """Press button."""
         try:
-            await self.entity_description.press_fn(self.coordinator.device)
+            await self.entity_description.press_fn(self.coordinator)
         except RequestNotSuccessful as exc:
             raise HomeAssistantError(
                 translation_domain=DOMAIN,
@@ -68,4 +81,3 @@ class LaMarzoccoButtonEntity(LaMarzoccoEntity, ButtonEntity):
                     "key": self.entity_description.key,
                 },
             ) from exc
-        await self.coordinator.async_request_refresh()

tests/components/lamarzocco/test_button.py

@@ -1,6 +1,6 @@
 """Tests for the La Marzocco Buttons."""
 
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from pylamarzocco.exceptions import RequestNotSuccessful
 import pytest
@@ -33,14 +33,18 @@ async def test_start_backflush(
     assert entry
     assert entry == snapshot
 
-    await hass.services.async_call(
-        BUTTON_DOMAIN,
-        SERVICE_PRESS,
-        {
-            ATTR_ENTITY_ID: f"button.{serial_number}_start_backflush",
-        },
-        blocking=True,
-    )
+    with patch(
+        "homeassistant.components.lamarzocco.button.asyncio.sleep",
+        new_callable=AsyncMock,
+    ):
+        await hass.services.async_call(
+            BUTTON_DOMAIN,
+            SERVICE_PRESS,
+            {
+                ATTR_ENTITY_ID: f"button.{serial_number}_start_backflush",
+            },
+            blocking=True,
+        )
 
     assert len(mock_lamarzocco.start_backflush.mock_calls) == 1
     mock_lamarzocco.start_backflush.assert_called_once()