Fix KeyError in hydrawise (#169853)

homeassistant/components/hydrawise/__init__.py

@@ -46,6 +46,11 @@ async def async_setup_entry(
     water_use_coordinator = HydrawiseWaterUseDataUpdateCoordinator(
         hass, config_entry, hydrawise, main_coordinator
     )
+    # async_track_zones is registered first on water_use_coordinator,
+    # so the water-use coordinator's data is in sync before
+    # callbacks below construct entities for newly added zones.
+    water_use_coordinator.async_track_zones()
+    main_coordinator.async_track_zones()
     await water_use_coordinator.async_config_entry_first_refresh()
     config_entry.runtime_data = HydrawiseUpdateCoordinators(
         main=main_coordinator,

homeassistant/components/hydrawise/coordinator.py

@@ -82,6 +82,10 @@ class HydrawiseMainDataUpdateCoordinator(HydrawiseDataUpdateCoordinator):
         self.new_zones_callbacks: list[
             Callable[[Iterable[tuple[Zone, Controller]]], None]
         ] = []
+
+    @callback
+    def async_track_zones(self) -> None:
+        """Begin tracking zone and controller add/remove on updates."""
         self.async_add_listener(self._add_remove_zones)
 
     async def _async_update_data(self) -> HydrawiseData:
@@ -198,6 +202,23 @@ class HydrawiseWaterUseDataUpdateCoordinator(HydrawiseDataUpdateCoordinator):
         self.api = api
         self._main_coordinator = main_coordinator
 
+    @callback
+    def async_track_zones(self) -> None:
+        """Begin tracking zone and controller add/remove on updates."""
+        self._main_coordinator.async_add_listener(self._sync_data_from_main)
+
+    @callback
+    def _sync_data_from_main(self) -> None:
+        """Sync data references from the main coordinator after it updates."""
+        if self.data is None or self._main_coordinator.data is None:
+            return  # type: ignore[unreachable]
+        main_data = self._main_coordinator.data
+        self.data.user = main_data.user
+        self.data.controllers = main_data.controllers
+        self.data.zones = main_data.zones
+        self.data.zone_id_to_controller = main_data.zone_id_to_controller
+        self.data.sensors = main_data.sensors
+
     async def _async_update_data(self) -> HydrawiseData:
         """Fetch the latest data from Hydrawise."""
         daily_water_summary: dict[int, ControllerWaterUseSummary] = {}

tests/components/hydrawise/test_init.py

@@ -66,6 +66,12 @@ async def test_auto_add_devices(
     # 1 controller + 2 zones
     assert len(all_devices) == 3
 
+    # Sensors that use the water-use coordinator should exist for the initial zones.
+    assert hass.states.get("sensor.zone_one_daily_active_water_use") is not None
+    assert hass.states.get("sensor.zone_one_daily_active_watering_time") is not None
+    assert hass.states.get("sensor.zone_two_daily_active_water_use") is not None
+    assert hass.states.get("sensor.zone_two_daily_active_watering_time") is not None
+
     controller2 = deepcopy(controller)
     controller2.id += 10
     controller2.name += " 2"
@@ -100,6 +106,11 @@ async def test_auto_add_devices(
     # 2 controllers + 4 zones
     assert len(all_devices) == 6
 
+    # Sensors that use the water-use coordinator must also be created for the
+    # newly added zones, even though the water-use coordinator hasn't refreshed.
+    assert hass.states.get("sensor.zone_one_2_daily_active_watering_time") is not None
+    assert hass.states.get("sensor.zone_two_2_daily_active_watering_time") is not None
+
 
 async def test_auto_remove_devices(
     hass: HomeAssistant,