Add yeelight service to enable disable music mode (#44533)

* Add service to enable / disable music mode

* Black reformat

* Update test

* Fix tests

* Revert consts cleanup

* Use entity method as service call

* Use ATTR for service call

* Sort

* Add tests

* Fix isort

* Fix print

* Black

homeassistant/components/yeelight/__init__.py

@@ -53,6 +53,7 @@ DATA_UNSUB_UPDATE_LISTENER = "unsub_update_listener"
 ATTR_COUNT = "count"
 ATTR_ACTION = "action"
 ATTR_TRANSITIONS = "transitions"
+ATTR_MODE_MUSIC = "music_mode"
 
 ACTION_RECOVER = "recover"
 ACTION_STAY = "stay"

homeassistant/components/yeelight/light.py

@@ -49,6 +49,7 @@ from . import (
     ACTION_RECOVER,
     ATTR_ACTION,
     ATTR_COUNT,
+    ATTR_MODE_MUSIC,
     ATTR_TRANSITIONS,
     CONF_FLOW_PARAMS,
     CONF_MODE_MUSIC,
@@ -77,6 +78,7 @@ SUPPORT_YEELIGHT_RGB = SUPPORT_YEELIGHT_WHITE_TEMP | SUPPORT_COLOR
 ATTR_MINUTES = "minutes"
 
 SERVICE_SET_MODE = "set_mode"
+SERVICE_SET_MUSIC_MODE = "set_music_mode"
 SERVICE_START_FLOW = "start_flow"
 SERVICE_SET_COLOR_SCENE = "set_color_scene"
 SERVICE_SET_HSV_SCENE = "set_hsv_scene"
@@ -175,6 +177,10 @@ SERVICE_SCHEMA_SET_MODE = {
     vol.Required(ATTR_MODE): vol.In([mode.name.lower() for mode in PowerMode])
 }
 
+SERVICE_SCHEMA_SET_MUSIC_MODE = {
+    vol.Required(ATTR_MODE_MUSIC): cv.boolean,
+}
+
 SERVICE_SCHEMA_START_FLOW = YEELIGHT_FLOW_TRANSITION_SCHEMA
 
 SERVICE_SCHEMA_SET_COLOR_SCENE = {
@@ -404,6 +410,11 @@ def _async_setup_services(hass: HomeAssistant):
         SERVICE_SCHEMA_SET_AUTO_DELAY_OFF_SCENE,
         _async_set_auto_delay_off_scene,
     )
+    platform.async_register_entity_service(
+        SERVICE_SET_MUSIC_MODE,
+        SERVICE_SCHEMA_SET_MUSIC_MODE,
+        "set_music_mode",
+    )
 
 
 class YeelightGenericLight(YeelightEntity, LightEntity):
@@ -550,7 +561,11 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
     def device_state_attributes(self):
         """Return the device specific state attributes."""
 
-        attributes = {"flowing": self.device.is_color_flow_enabled}
+        attributes = {
+            "flowing": self.device.is_color_flow_enabled,
+            "music_mode": self._bulb.music_mode,
+        }
+
         if self.device.is_nightlight_supported:
             attributes["night_light"] = self.device.is_nightlight_enabled
 
@@ -591,13 +606,18 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
 
         return color_util.color_RGB_to_hs(red, green, blue)
 
-    def set_music_mode(self, mode) -> None:
+    def set_music_mode(self, music_mode) -> None:
         """Set the music mode on or off."""
-        if mode:
-            self._bulb.start_music()
+        if music_mode:
+            try:
+                self._bulb.start_music()
+            except AssertionError as ex:
+                _LOGGER.error(ex)
         else:
             self._bulb.stop_music()
 
+        self.device.update()
+
     @_cmd
     def set_brightness(self, brightness, duration) -> None:
         """Set bulb brightness."""

homeassistant/components/yeelight/services.yaml

@@ -85,3 +85,12 @@ start_flow:
     transitions:
       description: Array of transitions, for desired effect. Examples https://yeelight.readthedocs.io/en/stable/flow.html
       example: '[{ "TemperatureTransition": [1900, 1000, 80] }, { "TemperatureTransition": [1900, 1000, 10] }]'
+set_music_mode:
+  description: Enable or disable music_mode
+  fields:
+    entity_id:
+      description: Name of the light entity.
+      example: "light.yeelight"
+    music_mode:
+      description: Use true or false to enable / disable music_mode
+      example: true

tests/components/yeelight/test_light.py

@@ -31,6 +31,7 @@ from homeassistant.components.light import (
 )
 from homeassistant.components.yeelight import (
     ATTR_COUNT,
+    ATTR_MODE_MUSIC,
     ATTR_TRANSITIONS,
     CONF_CUSTOM_EFFECTS,
     CONF_FLOW_PARAMS,
@@ -72,6 +73,7 @@ from homeassistant.components.yeelight.light import (
     SERVICE_SET_COLOR_TEMP_SCENE,
     SERVICE_SET_HSV_SCENE,
     SERVICE_SET_MODE,
+    SERVICE_SET_MUSIC_MODE,
     SERVICE_START_FLOW,
     SUPPORT_YEELIGHT,
     SUPPORT_YEELIGHT_RGB,
@@ -135,7 +137,14 @@ async def test_services(hass: HomeAssistant, caplog):
         assert await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
-    async def _async_test_service(service, data, method, payload=None, domain=DOMAIN):
+    async def _async_test_service(
+        service,
+        data,
+        method,
+        payload=None,
+        domain=DOMAIN,
+        failure_side_effect=BulbException,
+    ):
         err_count = len([x for x in caplog.records if x.levelno == logging.ERROR])
 
         # success
@@ -153,13 +162,14 @@ async def test_services(hass: HomeAssistant, caplog):
         )
 
         # failure
-        mocked_method = MagicMock(side_effect=BulbException)
-        setattr(type(mocked_bulb), method, mocked_method)
-        await hass.services.async_call(domain, service, data, blocking=True)
-        assert (
-            len([x for x in caplog.records if x.levelno == logging.ERROR])
-            == err_count + 1
-        )
+        if failure_side_effect:
+            mocked_method = MagicMock(side_effect=failure_side_effect)
+            setattr(type(mocked_bulb), method, mocked_method)
+            await hass.services.async_call(domain, service, data, blocking=True)
+            assert (
+                len([x for x in caplog.records if x.levelno == logging.ERROR])
+                == err_count + 1
+            )
 
     # turn_on
     brightness = 100
@@ -283,6 +293,29 @@ async def test_services(hass: HomeAssistant, caplog):
         [SceneClass.AUTO_DELAY_OFF, 50, 1],
     )
 
+    # set_music_mode failure enable
+    await _async_test_service(
+        SERVICE_SET_MUSIC_MODE,
+        {ATTR_ENTITY_ID: ENTITY_LIGHT, ATTR_MODE_MUSIC: "true"},
+        "start_music",
+        failure_side_effect=AssertionError,
+    )
+
+    # set_music_mode disable
+    await _async_test_service(
+        SERVICE_SET_MUSIC_MODE,
+        {ATTR_ENTITY_ID: ENTITY_LIGHT, ATTR_MODE_MUSIC: "false"},
+        "stop_music",
+        failure_side_effect=None,
+    )
+
+    # set_music_mode success enable
+    await _async_test_service(
+        SERVICE_SET_MUSIC_MODE,
+        {ATTR_ENTITY_ID: ENTITY_LIGHT, ATTR_MODE_MUSIC: "true"},
+        "start_music",
+        failure_side_effect=None,
+    )
     # test _cmd wrapper error handler
     err_count = len([x for x in caplog.records if x.levelno == logging.ERROR])
     type(mocked_bulb).turn_on = MagicMock()
@@ -338,6 +371,7 @@ async def test_device_types(hass: HomeAssistant):
         target_properties["friendly_name"] = name
         target_properties["flowing"] = False
         target_properties["night_light"] = True
+        target_properties["music_mode"] = False
         assert dict(state.attributes) == target_properties
 
         await hass.config_entries.async_unload(config_entry.entry_id)
@@ -365,6 +399,7 @@ async def test_device_types(hass: HomeAssistant):
         nightlight_properties["icon"] = "mdi:weather-night"
         nightlight_properties["flowing"] = False
         nightlight_properties["night_light"] = True
+        nightlight_properties["music_mode"] = False
         assert dict(state.attributes) == nightlight_properties
 
         await hass.config_entries.async_unload(config_entry.entry_id)