Add event entity on websocket ready in Husqvarna Automower (home-assistant#151428)

Author: Thomas55555

homeassistant/components/husqvarna_automower/event.py

@@ -36,27 +36,34 @@ async def async_setup_entry(
     """Set up Automower message event entities.
 
     Entities are created dynamically based on messages received from the API,
-    but only for mowers that support message events.
+    but only for mowers that support message events after the WebSocket connection
+    is ready.
     """
     coordinator = config_entry.runtime_data
     entity_registry = er.async_get(hass)
 
-    restored_mowers = {
+    restored_mowers: set[str] = {
         entry.unique_id.removesuffix("_message")
         for entry in er.async_entries_for_config_entry(
             entity_registry, config_entry.entry_id
         )
         if entry.domain == EVENT_DOMAIN
     }
 
-    async_add_entities(
-        AutomowerMessageEventEntity(mower_id, coordinator)
-        for mower_id in restored_mowers
-        if mower_id in coordinator.data
-    )
+    @callback
+    def _on_ws_ready() -> None:
+        async_add_entities(
+            AutomowerMessageEventEntity(mower_id, coordinator, websocket_alive=True)
+            for mower_id in restored_mowers
+            if mower_id in coordinator.data
+        )
+        coordinator.api.unregister_ws_ready_callback(_on_ws_ready)
+
+    coordinator.api.register_ws_ready_callback(_on_ws_ready)
 
     @callback
     def _handle_message(msg: SingleMessageData) -> None:
+        """Add entity dynamically if a new mower sends messages."""
         if msg.id in restored_mowers:
             return
 
@@ -78,11 +85,17 @@ def __init__(
         self,
         mower_id: str,
         coordinator: AutomowerDataUpdateCoordinator,
+        *,
+        websocket_alive: bool | None = None,
     ) -> None:
         """Initialize Automower message event entity."""
         super().__init__(mower_id, coordinator)
         self._attr_unique_id = f"{mower_id}_message"
-        self.websocket_alive: bool = coordinator.websocket_alive
+        self.websocket_alive: bool = (
+            websocket_alive
+            if websocket_alive is not None
+            else coordinator.websocket_alive
+        )
 
     @property
     def available(self) -> bool: