Improve color_mode handling when resetting lights

- Add support for color_temp, xy and rgb color modes

apps/media_lights_sync/media_lights_sync.py

@@ -13,6 +13,11 @@
 
 PICTURE_ATTRIBUTES = ["entity_picture_local", "entity_picture"]
 
+COLOR_MODES = {
+    "xy": "xy_color",
+    "color_temp": "color_temp",
+}
+
 
 class MediaLightsSync(hass.Hass):
     """MediaLightsSync class."""
@@ -84,14 +89,14 @@ def reset_lights(self):
             for i in range(len(self.lights)):
                 state = self.initial_lights_states[i]["state"]
                 attributes = self.initial_lights_states[i]["attributes"]
-                color_attr = "color_temp" if attributes.get("color_mode", "rgb") == "color_temp" else "rgb_color"
+                color_attr = COLOR_MODES.get(attributes.get("color_mode", None), "rgb_color")
 
-                self.set_light(state.lower(), self.lights[i], color_attr=color_attr, color=attributes.get(color_attr, None),
+                self.set_light(state.lower(), self.lights[i], color=attributes.get(color_attr, None), color_attr=color_attr,
                                brightness=attributes.get("brightness", None), transition=self.transition)
             self.initial_lights_states = None
             self.media_player_callbacks = {}
 
-    def set_light(self, new_state, entity, color_attr="rgb_color", color=None, brightness=None, transition=None):
+    def set_light(self, new_state, entity, color=None, color_attr="rgb_color", brightness=None, transition=None):
         """Change the color of a light."""
         attributes = {}
         if transition is not None:

tests/media_lights_sync_test.py

@@ -189,6 +189,20 @@ def test_can_change_lights_after_reset(self, assert_that, media_player, given_th
         assert_that('light.test_light_1').was.turned_on(brightness=255, rgb_color=[59, 180, 180])
         assert_that('light.test_light_2').was.turned_on(brightness=255, rgb_color=[46, 56, 110])
 
+    def test_can_reset_color_modes(self, assert_that, media_player, given_that, update_passed_args, hass_mocks):
+        with update_passed_args():
+            given_that.passed_arg('reset_lights_after').is_set_to(True)
+        given_that.state_of('light.test_light_1').is_set_to('on', {'color_mode': 'color_temp', 'color_temp': 500, 'brightness': 150})
+        given_that.state_of('light.test_light_2').is_set_to('on', {'color_mode': 'xy', 'xy_color': [0.166, 0.269], 'brightness': 100})
+
+        media_player('media_player.tv_test').update_state('playing', {"entity_picture": rgb_images[0]})
+        given_that.mock_functions_are_cleared()
+
+        media_player('media_player.tv_test').update_state('idle')
+
+        assert_that('light.test_light_1').was.turned_on(color_temp=500, brightness=150)
+        assert_that('light.test_light_2').was.turned_on(xy_color=[0.166, 0.269], brightness=100)
+
 
 class TestURLErrors:
     def test_url_error_are_handled(self, media_player, hass_errors, hass_mocks):