Use entry.async_on_unload in fireservicerota (#138360)

home-assistant · Feb 12, 2025 · 2033dbd · 2033dbd
1 parent a6c5144
commit 2033dbd
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/homeassistant/components/fireservicerota/__init__.py b/homeassistant/components/fireservicerota/__init__.py
@@ -27,6 +27,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     if client.token_refresh_failure:
         return False
 
+    entry.async_on_unload(client.async_stop_listener)
     coordinator = FireServiceUpdateCoordinator(hass, client, entry)
 
     await coordinator.async_config_entry_first_refresh()
@@ -43,10 +44,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload FireServiceRota config entry."""
-
-    await hass.async_add_executor_job(
-        hass.data[DOMAIN][entry.entry_id][DATA_CLIENT].websocket.stop_listener
-    )
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
     if unload_ok:
         del hass.data[DOMAIN][entry.entry_id]

diff --git a/homeassistant/components/fireservicerota/coordinator.py b/homeassistant/components/fireservicerota/coordinator.py
@@ -213,3 +213,7 @@ async def async_set_response(self, value) -> None:
         )
 
         await self.update_call(self.fsr.set_incident_response, self.incident_id, value)
+
+    async def async_stop_listener(self) -> None:
+        """Stop listener."""
+        await self._hass.async_add_executor_job(self.websocket.stop_listener)