Improve color mode handling in MockLight (#109298)

tests/testing_config/custom_components/test/light.py

@@ -2,7 +2,7 @@
 
 Call init before using it in your tests to ensure clean test data.
 """
-from homeassistant.components.light import LightEntity
+from homeassistant.components.light import ColorMode, LightEntity
 from homeassistant.const import STATE_OFF, STATE_ON
 
 from tests.common import MockToggleEntity
@@ -32,13 +32,21 @@ async def async_setup_platform(
     async_add_entities_callback(ENTITIES)
 
 
+TURN_ON_ARG_TO_COLOR_MODE = {
+    "hs_color": ColorMode.HS,
+    "xy_color": ColorMode.XY,
+    "rgb_color": ColorMode.RGB,
+    "rgbw_color": ColorMode.RGBW,
+    "rgbww_color": ColorMode.RGBWW,
+    "color_temp_kelvin": ColorMode.COLOR_TEMP,
+}
+
+
 class MockLight(MockToggleEntity, LightEntity):
     """Mock light class."""
 
-    color_mode = None
     _attr_max_color_temp_kelvin = 6500
     _attr_min_color_temp_kelvin = 2000
-    supported_color_modes = None
     supported_features = 0
 
     brightness = None
@@ -49,6 +57,23 @@ class MockLight(MockToggleEntity, LightEntity):
     rgbww_color = None
     xy_color = None
 
+    def __init__(
+        self,
+        name,
+        state,
+        unique_id=None,
+        supported_color_modes: set[ColorMode] | None = None,
+    ):
+        """Initialize the mock light."""
+        super().__init__(name, state, unique_id)
+        if supported_color_modes is None:
+            supported_color_modes = {ColorMode.ONOFF}
+        self._attr_supported_color_modes = supported_color_modes
+        color_mode = ColorMode.UNKNOWN
+        if len(supported_color_modes) == 1:
+            color_mode = next(iter(supported_color_modes))
+        self._attr_color_mode = color_mode
+
     def turn_on(self, **kwargs):
         """Turn the entity on."""
         super().turn_on(**kwargs)
@@ -65,3 +90,5 @@ class MockLight(MockToggleEntity, LightEntity):
                 setattr(self, key, value)
             if key == "white":
                 setattr(self, "brightness", value)
+            if key in TURN_ON_ARG_TO_COLOR_MODE:
+                self._attr_color_mode = TURN_ON_ARG_TO_COLOR_MODE[key]