From 52c773a776500033c3fe4c286133fd1e07ddf7b9 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 11 May 2023 19:02:32 +0900
Subject: [PATCH] Always request at least one zone for multi-zone LIFX devices
 (#92683)

---
 homeassistant/components/lifx/coordinator.py | 47 +++++++++++++++++---
 tests/components/lifx/test_light.py          | 10 +++++
 2 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/lifx/coordinator.py b/homeassistant/components/lifx/coordinator.py
index f8964e78f63..e668a7ad79a 100644
--- a/homeassistant/components/lifx/coordinator.py
+++ b/homeassistant/components/lifx/coordinator.py
@@ -11,6 +11,7 @@ from typing import Any, cast
 
 from aiolifx.aiolifx import (
     Light,
+    Message,
     MultiZoneDirection,
     MultiZoneEffectType,
     TileEffectType,
@@ -56,6 +57,8 @@ from .util import (
 LIGHT_UPDATE_INTERVAL = 10
 REQUEST_REFRESH_DELAY = 0.35
 LIFX_IDENTIFY_DELAY = 3.0
+ZONES_PER_COLOR_UPDATE_REQUEST = 8
+
 RSSI_DBM_FW = AwesomeVersion("2.77")
 
 
@@ -208,18 +211,50 @@ class LIFXUpdateCoordinator(DataUpdateCoordinator[None]):
     def get_number_of_zones(self) -> int:
         """Return the number of zones.
 
-        If the number of zones is not yet populated, return 0
+        If the number of zones is not yet populated, return 1 since
+        the device will have a least one zone.
         """
-        return len(self.device.color_zones) if self.device.color_zones else 0
+        return len(self.device.color_zones) if self.device.color_zones else 1
 
     @callback
     def _async_build_color_zones_update_requests(self) -> list[Callable]:
         """Build a color zones update request."""
         device = self.device
-        return [
-            partial(device.get_color_zones, start_index=zone)
-            for zone in range(0, self.get_number_of_zones(), 8)
-        ]
+        calls: list[Callable] = []
+        for zone in range(
+            0, self.get_number_of_zones(), ZONES_PER_COLOR_UPDATE_REQUEST
+        ):
+
+            def _wrap_get_color_zones(
+                callb: Callable[[Message, dict[str, Any] | None], None],
+                get_color_zones_args: dict[str, Any],
+            ) -> None:
+                """Capture the callback and make sure resp_set_multizonemultizone is called before."""
+
+                def _wrapped_callback(
+                    bulb: Light,
+                    response: Message,
+                    **kwargs: Any,
+                ) -> None:
+                    # We need to call resp_set_multizonemultizone to populate
+                    # the color_zones attribute before calling the callback
+                    device.resp_set_multizonemultizone(response)
+                    # Now call the original callback
+                    callb(bulb, response, **kwargs)
+
+                device.get_color_zones(**get_color_zones_args, callb=_wrapped_callback)
+
+            calls.append(
+                partial(
+                    _wrap_get_color_zones,
+                    get_color_zones_args={
+                        "start_index": zone,
+                        "end_index": zone + ZONES_PER_COLOR_UPDATE_REQUEST - 1,
+                    },
+                )
+            )
+
+        return calls
 
     async def _async_update_data(self) -> None:
         """Fetch all device data from the api."""
diff --git a/tests/components/lifx/test_light.py b/tests/components/lifx/test_light.py
index fe68bd6547a..67bf2754d11 100644
--- a/tests/components/lifx/test_light.py
+++ b/tests/components/lifx/test_light.py
@@ -1754,6 +1754,8 @@ async def test_light_strip_zones_not_populated_yet(hass: HomeAssistant) -> None:
     bulb.power_level = 65535
     bulb.color_zones = None
     bulb.color = [65535, 65535, 65535, 65535]
+    assert bulb.get_color_zones.calls == []
+
     with _patch_discovery(device=bulb), _patch_config_flow_try_connect(
         device=bulb
     ), _patch_device(device=bulb):
@@ -1761,6 +1763,14 @@ async def test_light_strip_zones_not_populated_yet(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
 
     entity_id = "light.my_bulb"
+    # Make sure we at least try to fetch the first zone
+    # to ensure we populate the zones from the 503 response
+    assert len(bulb.get_color_zones.calls) == 3
+    # Once to populate the number of zones
+    assert bulb.get_color_zones.calls[0][1]["start_index"] == 0
+    # Again once we know the number of zones
+    assert bulb.get_color_zones.calls[1][1]["start_index"] == 0
+    assert bulb.get_color_zones.calls[2][1]["start_index"] == 8
 
     state = hass.states.get(entity_id)
     assert state.state == "on"