diff --git a/homeassistant/components/lamarzocco/binary_sensor.py b/homeassistant/components/lamarzocco/binary_sensor.py index 98cf7cf222e..9bf04129095 100644 --- a/homeassistant/components/lamarzocco/binary_sensor.py +++ b/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( diff --git a/homeassistant/components/lamarzocco/coordinator.py b/homeassistant/components/lamarzocco/coordinator.py index 751ef550516..f0f64e02c28 100644 --- a/homeassistant/components/lamarzocco/coordinator.py +++ b/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( - update_callback=lambda _: self.async_set_updated_data(None), - connect_callback=self.async_update_listeners, - disconnect_callback=self.async_update_listeners, - ), + target=self.connect_websocket(), 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.""" diff --git a/homeassistant/components/lamarzocco/entity.py b/homeassistant/components/lamarzocco/entity.py index 2e3a7f2ce83..6dc024645ce 100644 --- a/homeassistant/components/lamarzocco/entity.py +++ b/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__( diff --git a/homeassistant/components/lamarzocco/number.py b/homeassistant/components/lamarzocco/number.py index 81a03b4d6ee..7c4fe33a041 100644 --- a/homeassistant/components/lamarzocco/number.py +++ b/homeassistant/components/lamarzocco/number.py @@ -100,8 +100,9 @@ ENTITIES: tuple[LaMarzoccoNumberEntityDescription, ...] = ( .seconds.seconds_out ), available_fn=( - lambda machine: cast( - PreBrewing, machine.dashboard.config[WidgetType.CM_PRE_BREWING] + lambda coordinator: cast( + 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( - PreBrewing, machine.dashboard.config[WidgetType.CM_PRE_BREWING] + available_fn=lambda coordinator: cast( + 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( - PreBrewing, machine.dashboard.config[WidgetType.CM_PRE_BREWING] + lambda coordinator: cast( + PreBrewing, + coordinator.device.dashboard.config[WidgetType.CM_PRE_BREWING], ).mode is PreExtractionMode.PREBREWING ), diff --git a/tests/components/lamarzocco/test_binary_sensor.py b/tests/components/lamarzocco/test_binary_sensor.py index 8e92c9bbba9..570b5aef8ec 100644 --- a/tests/components/lamarzocco/test_binary_sensor.py +++ b/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,