Fix intermittent unavailability for lamarzocco brew active sensor (#144120)

* Fix brew active intermittent unavailability for lamarzocco

* Whitespaces

homeassistant/components/lamarzocco/binary_sensor.py

@@ -52,7 +52,7 @@ ENTITIES: tuple[LaMarzoccoBinarySensorEntityDescription, ...] = (
             ).status
             is MachineState.BREWING
         ),
-        available_fn=lambda device: device.websocket.connected,
+        available_fn=lambda coordinator: not coordinator.websocket_terminated,
         entity_category=EntityCategory.DIAGNOSTIC,
     ),
     LaMarzoccoBinarySensorEntityDescription(

homeassistant/components/lamarzocco/coordinator.py

@@ -44,6 +44,7 @@ class LaMarzoccoUpdateCoordinator(DataUpdateCoordinator[None]):
 
     _default_update_interval = SCAN_INTERVAL
     config_entry: LaMarzoccoConfigEntry
+    websocket_terminated = True
 
     def __init__(
         self,
@@ -92,15 +93,9 @@ class LaMarzoccoConfigUpdateCoordinator(LaMarzoccoUpdateCoordinator):
         await self.device.get_dashboard()
         _LOGGER.debug("Current status: %s", self.device.dashboard.to_dict())
 
-        _LOGGER.debug("Init WebSocket in background task")
-
         self.config_entry.async_create_background_task(
             hass=self.hass,
-            target=self.device.connect_dashboard_websocket(
+            target=self.connect_websocket(),
-                update_callback=lambda _: self.async_set_updated_data(None),
-                connect_callback=self.async_update_listeners,
-                disconnect_callback=self.async_update_listeners,
-            ),
             name="lm_websocket_task",
         )
 
@@ -112,6 +107,23 @@ class LaMarzoccoConfigUpdateCoordinator(LaMarzoccoUpdateCoordinator):
         )
         self.config_entry.async_on_unload(websocket_close)
 
+    async def connect_websocket(self) -> None:
+        """Connect to the websocket."""
+
+        _LOGGER.debug("Init WebSocket in background task")
+
+        self.websocket_terminated = False
+        self.async_update_listeners()
+
+        await self.device.connect_dashboard_websocket(
+            update_callback=lambda _: self.async_set_updated_data(None),
+            connect_callback=self.async_update_listeners,
+            disconnect_callback=self.async_update_listeners,
+        )
+
+        self.websocket_terminated = True
+        self.async_update_listeners()
+
 
 class LaMarzoccoSettingsUpdateCoordinator(LaMarzoccoUpdateCoordinator):
     """Coordinator for La Marzocco settings."""

homeassistant/components/lamarzocco/entity.py

@@ -3,7 +3,6 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 
-from pylamarzocco import LaMarzoccoMachine
 from pylamarzocco.const import FirmwareType
 
 from homeassistant.const import CONF_ADDRESS, CONF_MAC
@@ -23,7 +22,7 @@ from .coordinator import LaMarzoccoUpdateCoordinator
 class LaMarzoccoEntityDescription(EntityDescription):
     """Description for all LM entities."""
 
-    available_fn: Callable[[LaMarzoccoMachine], bool] = lambda _: True
+    available_fn: Callable[[LaMarzoccoUpdateCoordinator], bool] = lambda _: True
     supported_fn: Callable[[LaMarzoccoUpdateCoordinator], bool] = lambda _: True
 
 
@@ -74,7 +73,7 @@ class LaMarzoccoEntity(LaMarzoccoBaseEntity):
     def available(self) -> bool:
         """Return True if entity is available."""
         if super().available:
-            return self.entity_description.available_fn(self.coordinator.device)
+            return self.entity_description.available_fn(self.coordinator)
         return False
 
     def __init__(

homeassistant/components/lamarzocco/number.py

@@ -100,8 +100,9 @@ ENTITIES: tuple[LaMarzoccoNumberEntityDescription, ...] = (
             .seconds.seconds_out
         ),
         available_fn=(
-            lambda machine: cast(
+            lambda coordinator: cast(
-                PreBrewing, machine.dashboard.config[WidgetType.CM_PRE_BREWING]
+                PreBrewing,
+                coordinator.device.dashboard.config[WidgetType.CM_PRE_BREWING],
             ).mode
             is PreExtractionMode.PREINFUSION
         ),
@@ -140,8 +141,8 @@ ENTITIES: tuple[LaMarzoccoNumberEntityDescription, ...] = (
             .times.pre_brewing[0]
             .seconds.seconds_in
         ),
-        available_fn=lambda machine: cast(
+        available_fn=lambda coordinator: cast(
-            PreBrewing, machine.dashboard.config[WidgetType.CM_PRE_BREWING]
+            PreBrewing, coordinator.device.dashboard.config[WidgetType.CM_PRE_BREWING]
         ).mode
         is PreExtractionMode.PREBREWING,
         supported_fn=(
@@ -180,8 +181,9 @@ ENTITIES: tuple[LaMarzoccoNumberEntityDescription, ...] = (
             .seconds.seconds_out
         ),
         available_fn=(
-            lambda machine: cast(
+            lambda coordinator: cast(
-                PreBrewing, machine.dashboard.config[WidgetType.CM_PRE_BREWING]
+                PreBrewing,
+                coordinator.device.dashboard.config[WidgetType.CM_PRE_BREWING],
             ).mode
             is PreExtractionMode.PREBREWING
         ),

tests/components/lamarzocco/test_binary_sensor.py

@@ -1,5 +1,6 @@
 """Tests for La Marzocco binary sensors."""
 
+from collections.abc import Generator
 from datetime import timedelta
 from unittest.mock import MagicMock, patch
 
@@ -33,6 +34,16 @@ async def test_binary_sensors(
     await snapshot_platform(hass, entity_registry, snapshot, mock_config_entry.entry_id)
 
 
+@pytest.fixture(autouse=True)
+def mock_websocket_terminated() -> Generator[bool]:
+    """Mock websocket terminated."""
+    with patch(
+        "homeassistant.components.lamarzocco.coordinator.LaMarzoccoUpdateCoordinator.websocket_terminated",
+        new=False,
+    ) as mock_websocket_terminated:
+        yield mock_websocket_terminated
+
+
 async def test_brew_active_unavailable(
     hass: HomeAssistant,
     mock_lamarzocco: MagicMock,