From f4d872a54089a05f84292e0f58cb8e9d95562181 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <basnijholt@gmail.com>
Date: Sun, 30 Jul 2023 13:01:08 -0700
Subject: [PATCH] Make sure that SimpleSwitches are added before AdaptiveSwitch
 (#686)

---
 custom_components/adaptive_lighting/switch.py | 24 +++++++++++++++----
 tests/test_switch.py                          |  3 +--
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/custom_components/adaptive_lighting/switch.py b/custom_components/adaptive_lighting/switch.py
index 58638c3d..2cf2c5db 100644
--- a/custom_components/adaptive_lighting/switch.py
+++ b/custom_components/adaptive_lighting/switch.py
@@ -293,7 +293,7 @@ def _switches_with_lights(
         entry = data.get(config.entry_id)
         if entry is None:  # entry might be disabled and therefore missing
             continue
-        switch = data[config.entry_id]["instance"]
+        switch = data[config.entry_id][SWITCH_DOMAIN]
         switch._expand_light_groups()
         # Check if any of the lights are in the switch's lights
         if set(switch.lights) & set(all_check_lights):
@@ -367,7 +367,7 @@ def _switches_from_service_call(
             ent_entry = ent_reg.async_get(entity_id)
             assert ent_entry is not None
             config_id = ent_entry.config_entry_id
-            switches.append(hass.data[DOMAIN][config_id]["instance"])
+            switches.append(hass.data[DOMAIN][config_id][SWITCH_DOMAIN])
         return switches
 
     if lights:
@@ -498,9 +498,6 @@ async def async_setup_entry(  # noqa: PLR0915
         adapt_brightness_switch,
     )
 
-    # save our switch instance, allows us to make switch's entity_id optional in service calls.
-    hass.data[DOMAIN][config_entry.entry_id]["instance"] = switch
-
     data[config_entry.entry_id][SLEEP_MODE_SWITCH] = sleep_mode_switch
     data[config_entry.entry_id][ADAPT_COLOR_SWITCH] = adapt_color_switch
     data[config_entry.entry_id][ADAPT_BRIGHTNESS_SWITCH] = adapt_brightness_switch
@@ -1016,6 +1013,23 @@ class AdaptiveSwitch(SwitchEntity, RestoreEntity):
             _LOGGER.debug("%s: Cancelled '_setup_listeners'", self._name)
             return
 
+        while not all(
+            sw._state is not None
+            for sw in [
+                self.sleep_mode_switch,
+                self.adapt_brightness_switch,
+                self.adapt_color_switch,
+            ]
+        ):
+            # Waits until `async_added_to_hass` is done, which in SimpleSwitch
+            # is when `_state` is set to `True` or `False`.
+            # Fixes first issue in https://github.com/basnijholt/adaptive-lighting/issues/682
+            _LOGGER.debug(
+                "%s: Waiting for simple switches to be initialized",
+                self._name,
+            )
+            await asyncio.sleep(0.1)
+
         assert not self.remove_listeners
 
         self._update_time_interval_listener()
diff --git a/tests/test_switch.py b/tests/test_switch.py
index be0c7d37..2a49f41c 100644
--- a/tests/test_switch.py
+++ b/tests/test_switch.py
@@ -331,9 +331,8 @@ async def test_adaptive_lighting_switches(hass):
     assert ADAPT_COLOR_SWITCH in data
     assert ADAPT_BRIGHTNESS_SWITCH in data
     assert UNDO_UPDATE_LISTENER in data
-    assert "instance" in data
 
-    assert len(data.keys()) == 6
+    assert len(data.keys()) == 5
 
 
 @pytest.mark.parametrize("lat,long,timezone", LAT_LONG_TZS)