From a38b51fcc55cb62f2970ef702175fcc50bff4bfe Mon Sep 17 00:00:00 2001
From: Bas Nijholt <bas@nijho.lt>
Date: Sun, 30 Jul 2023 00:04:37 -0700
Subject: [PATCH] More tests

---
 custom_components/adaptive_lighting/switch.py | 10 ++++-----
 tests/test_switch.py                          | 22 +++++++++++++++++--
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/custom_components/adaptive_lighting/switch.py b/custom_components/adaptive_lighting/switch.py
index 35501979..0a23a4dc 100644
--- a/custom_components/adaptive_lighting/switch.py
+++ b/custom_components/adaptive_lighting/switch.py
@@ -1509,11 +1509,6 @@ class AdaptiveSwitch(SwitchEntity, RestoreEntity):
                 event.context.id,
             )
 
-            if is_our_context(event.context):
-                # AFAIK this can only happen in _service_interceptor_turn_on_handler
-                # when calling _adapt_light
-                return
-
             if (
                 not self._detect_non_ha_changes
                 and not self.manager.is_proactively_adapting(event.context.id)
@@ -2080,10 +2075,13 @@ class AdaptiveLightingManager:
                 continue
 
             for eid in filtered_entity_ids:
+                context = switch.create_context("intercept")
+                self.clear_proactively_adapting(eid)
+                self.set_proactively_adapting(context.id, eid)
                 await switch._adapt_light(
                     light=eid,
                     # Must add a new context otherwise _adapt_light will bail out
-                    context=switch.create_context("intercept"),
+                    context=context,
                     transition=transition,
                 )
 
diff --git a/tests/test_switch.py b/tests/test_switch.py
index 33c53cbe..18dd96f5 100644
--- a/tests/test_switch.py
+++ b/tests/test_switch.py
@@ -1550,7 +1550,6 @@ async def setup_proactive_multiple_lights_two_switches(hass):
         {ATTR_ENTITY_ID: lights},
         blocking=True,
     )
-    assert all(hass.states.get(light).state == STATE_OFF for light in lights)
     defaults = {
         CONF_SUNRISE_TIME: datetime.time(SUNRISE.hour),
         CONF_SUNSET_TIME: datetime.time(SUNSET.hour),
@@ -1561,7 +1560,6 @@ async def setup_proactive_multiple_lights_two_switches(hass):
         CONF_MIN_COLOR_TEMP: 2500,  # to not coincide with sleep_color_temp}
         INTERNAL_CONF_PROACTIVE_SERVICE_CALL_ADAPTATION: True,
     }
-    assert all(hass.states.get(light) is not None for light in lights)
     _, switch1 = await setup_switch(
         hass, {CONF_NAME: "switch1", CONF_LIGHTS: [ENTITY_LIGHT_1], **defaults}
     )
@@ -1570,6 +1568,7 @@ async def setup_proactive_multiple_lights_two_switches(hass):
     )
     assert hass.states.get(switch1.entity_id).state == STATE_ON
     assert hass.states.get(switch2.entity_id).state == STATE_ON
+    assert all(hass.states.get(light).state == STATE_OFF for light in lights)
     return lights, switch1, switch2
 
 
@@ -1625,6 +1624,25 @@ async def test_proactive_multiple_lights_turn_on_non_managed_light(hass):
     assert len(turn_ons) == 1, turn_ons
 
 
+async def test_proactive_multiple_lights_turn_on_managed_lights_only(hass):
+    """Create switch and demo lights."""
+    lights, switch1, switch2 = await setup_proactive_multiple_lights_two_switches(hass)
+    _LOGGER.debug("Start test_proactive_multiple_lights_all_at_once")
+    # Setup demo lights and turn on
+    events = await _turn_on_and_track_event_contexts(
+        hass, "test1", lights[:-1], return_full_events=True
+    )
+    assert len(events) == 2, events
+
+    # Original turn_on call that is intercepted
+    assert events[0].context.id == "test1"
+    assert events[0].data["service_data"][ATTR_ENTITY_ID] == lights[:-1]
+
+    # The `has_intercepted` path
+    assert events[1].data["service_data"][ATTR_ENTITY_ID] == ENTITY_LIGHT_2
+    assert ":ntrc:" in events[1].context.id
+
+
 async def test_two_switches_for_single_light(hass):
     """Test the case where someone has two switches for a single light.