Lamarzocco fix websocket reconnect issue (home-assistant#156786)

Co-authored-by: copilot-swe-agent[bot] <[email protected]>

Author: 2 people

homeassistant/components/lamarzocco/coordinator.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from abc import abstractmethod
+from asyncio import Task
 from dataclasses import dataclass
 from datetime import timedelta
 import logging
@@ -44,7 +45,7 @@ class LaMarzoccoUpdateCoordinator(DataUpdateCoordinator[None]):
 
     _default_update_interval = SCAN_INTERVAL
     config_entry: LaMarzoccoConfigEntry
-    websocket_terminated = True
+    _websocket_task: Task | None = None
 
     def __init__(
         self,
@@ -64,6 +65,13 @@ def __init__(
         self.device = device
         self.cloud_client = cloud_client
 
+    @property
+    def websocket_terminated(self) -> bool:
+        """Return True if the websocket task is terminated or not running."""
+        if self._websocket_task is None:
+            return True
+        return self._websocket_task.done()
+
     async def _async_update_data(self) -> None:
         """Do the data update."""
         try:
@@ -95,13 +103,14 @@ async def _internal_async_update_data(self) -> None:
         # ensure token stays valid; does nothing if token is still valid
         await self.cloud_client.async_get_access_token()
 
-        if self.device.websocket.connected:
+        # Only skip websocket reconnection if it's currently connected and the task is still running
+        if self.device.websocket.connected and not self.websocket_terminated:
             return
 
         await self.device.get_dashboard()
         _LOGGER.debug("Current status: %s", self.device.dashboard.to_dict())
 
-        self.config_entry.async_create_background_task(
+        self._websocket_task = self.config_entry.async_create_background_task(
             hass=self.hass,
             target=self.connect_websocket(),
             name="lm_websocket_task",
@@ -120,7 +129,6 @@ async def connect_websocket(self) -> None:
 
         _LOGGER.debug("Init WebSocket in background task")
 
-        self.websocket_terminated = False
         self.async_update_listeners()
 
         await self.device.connect_dashboard_websocket(
@@ -129,7 +137,6 @@ async def connect_websocket(self) -> None:
             disconnect_callback=self.async_update_listeners,
         )
 
-        self.websocket_terminated = True
         self.async_update_listeners()
 
 

tests/components/lamarzocco/conftest.py

@@ -1,7 +1,7 @@
 """Lamarzocco session fixtures."""
 
 from collections.abc import Generator
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch
 
 from bleak.backends.device import BLEDevice
 from pylamarzocco.const import ModelName
@@ -132,6 +132,10 @@ def mock_lamarzocco(device_fixture: ModelName) -> Generator[MagicMock]:
             "schedule": machine_mock.schedule.to_dict(),
             "settings": machine_mock.settings.to_dict(),
         }
+        machine_mock.connect_dashboard_websocket = AsyncMock()
+        machine_mock.websocket = MagicMock()
+        machine_mock.websocket.connected = True
+        machine_mock.websocket.disconnect = AsyncMock()
         yield machine_mock
 
 
@@ -149,10 +153,11 @@ def mock_ble_device() -> BLEDevice:
 
 
 @pytest.fixture
-def mock_websocket_terminated() -> Generator[bool]:
+def mock_websocket_terminated() -> Generator[PropertyMock]:
     """Mock websocket terminated."""
     with patch(
         "homeassistant.components.lamarzocco.coordinator.LaMarzoccoUpdateCoordinator.websocket_terminated",
-        new=False,
+        new_callable=PropertyMock,
     ) as mock_websocket_terminated:
+        mock_websocket_terminated.return_value = False
         yield mock_websocket_terminated

tests/components/lamarzocco/test_binary_sensor.py

@@ -1,8 +1,7 @@
 """Tests for La Marzocco binary sensors."""
 
-from collections.abc import Generator
 from datetime import timedelta
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, PropertyMock, patch
 
 from freezegun.api import FrozenDateTimeFactory
 from pylamarzocco.exceptions import RequestNotSuccessful
@@ -36,24 +35,15 @@ 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,
     mock_config_entry: MockConfigEntry,
+    mock_websocket_terminated: PropertyMock,
 ) -> None:
     """Test the La Marzocco brew active becomes unavailable."""
 
-    mock_lamarzocco.websocket.connected = False
+    mock_websocket_terminated.return_value = True
     await async_init_integration(hass, mock_config_entry)
     state = hass.states.get(
         f"binary_sensor.{mock_lamarzocco.serial_number}_brewing_active"

tests/components/lamarzocco/test_init.py

@@ -1,7 +1,9 @@
 """Test initialization of lamarzocco."""
 
+from datetime import timedelta
 from unittest.mock import AsyncMock, MagicMock, patch
 
+from freezegun.api import FrozenDateTimeFactory
 from pylamarzocco.const import FirmwareType, ModelName
 from pylamarzocco.exceptions import AuthFail, RequestNotSuccessful
 from pylamarzocco.models import WebSocketDetails
@@ -30,7 +32,7 @@
     get_bluetooth_service_info,
 )
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 
 async def test_load_unload_config_entry(
@@ -310,3 +312,37 @@ async def test_device(
     device = device_registry.async_get(entry.device_id)
     assert device
     assert device == snapshot
+
+
+async def test_websocket_reconnects_after_termination(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_lamarzocco: MagicMock,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test the websocket reconnects after background task terminates."""
+    # Setup: websocket connected initially
+    mock_websocket = MagicMock()
+    mock_websocket.closed = False
+    mock_lamarzocco.websocket = WebSocketDetails(mock_websocket, None)
+
+    await async_init_integration(hass, mock_config_entry)
+
+    # Verify initial websocket connection was attempted
+    assert mock_lamarzocco.connect_dashboard_websocket.call_count == 1
+
+    # Simulate websocket disconnection (e.g., after internet outage)
+    mock_websocket.closed = True
+
+    # Simulate the background task terminating by patching websocket_terminated
+    with patch(
+        "homeassistant.components.lamarzocco.coordinator.LaMarzoccoConfigUpdateCoordinator.websocket_terminated",
+        new=True,
+    ):
+        # Trigger the coordinator's update (which runs every 60 seconds)
+        freezer.tick(timedelta(seconds=61))
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
+
+    # Verify websocket reconnection was attempted
+    assert mock_lamarzocco.connect_dashboard_websocket.call_count == 2