diff --git a/homeassistant/components/lamarzocco/coordinator.py b/homeassistant/components/lamarzocco/coordinator.py index 0b07409adb5..2385039f53d 100644 --- a/homeassistant/components/lamarzocco/coordinator.py +++ b/homeassistant/components/lamarzocco/coordinator.py @@ -91,9 +91,11 @@ class LaMarzoccoConfigUpdateCoordinator(LaMarzoccoUpdateCoordinator): _scale_address: str | None = None - async def _async_setup(self) -> None: + async def _async_connect_websocket(self) -> None: """Set up the coordinator.""" - if self._local_client is not None: + if self._local_client is not None and ( + self._local_client.websocket is None or self._local_client.websocket.closed + ): _LOGGER.debug("Init WebSocket in background task") self.config_entry.async_create_background_task( @@ -123,6 +125,7 @@ class LaMarzoccoConfigUpdateCoordinator(LaMarzoccoUpdateCoordinator): """Fetch data from API endpoint.""" await self.device.get_config() _LOGGER.debug("Current status: %s", str(self.device.config)) + await self._async_connect_websocket() self._async_add_remove_scale() @callback diff --git a/tests/components/lamarzocco/test_init.py b/tests/components/lamarzocco/test_init.py index 7d90c049a3b..fccfcbeef13 100644 --- a/tests/components/lamarzocco/test_init.py +++ b/tests/components/lamarzocco/test_init.py @@ -199,8 +199,11 @@ async def test_websocket_closed_on_unload( ) as local_client: client = local_client.return_value client.websocket = AsyncMock() - client.websocket.closed = False + await async_init_integration(hass, mock_config_entry) + mock_lamarzocco.websocket_connect.assert_called_once() + + client.websocket.closed = False hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP) await hass.async_block_till_done() client.websocket.close.assert_called_once()