MQTT Light: Use flash attribute in async_turn_off (#47919)

homeassistant/components/mqtt/light/schema_json.py

@@ -341,6 +341,18 @@ class MqttLightJson(MqttEntity, LightEntity, RestoreEntity):
         """Flag supported features."""
         return self._supported_features
 
+    def _set_flash_and_transition(self, message, **kwargs):
+        if ATTR_TRANSITION in kwargs:
+            message["transition"] = kwargs[ATTR_TRANSITION]
+
+        if ATTR_FLASH in kwargs:
+            flash = kwargs.get(ATTR_FLASH)
+
+            if flash == FLASH_LONG:
+                message["flash"] = self._flash_times[CONF_FLASH_TIME_LONG]
+            elif flash == FLASH_SHORT:
+                message["flash"] = self._flash_times[CONF_FLASH_TIME_SHORT]
+
     async def async_turn_on(self, **kwargs):
         """Turn the device on.
 
@@ -380,16 +392,7 @@ class MqttLightJson(MqttEntity, LightEntity, RestoreEntity):
                 self._hs = kwargs[ATTR_HS_COLOR]
                 should_update = True
 
-        if ATTR_FLASH in kwargs:
-            flash = kwargs.get(ATTR_FLASH)
-
-            if flash == FLASH_LONG:
-                message["flash"] = self._flash_times[CONF_FLASH_TIME_LONG]
-            elif flash == FLASH_SHORT:
-                message["flash"] = self._flash_times[CONF_FLASH_TIME_SHORT]
-
-        if ATTR_TRANSITION in kwargs:
-            message["transition"] = kwargs[ATTR_TRANSITION]
+        self._set_flash_and_transition(message, **kwargs)
 
         if ATTR_BRIGHTNESS in kwargs and self._config[CONF_BRIGHTNESS]:
             brightness_normalized = kwargs[ATTR_BRIGHTNESS] / DEFAULT_BRIGHTNESS_SCALE
@@ -449,8 +452,7 @@ class MqttLightJson(MqttEntity, LightEntity, RestoreEntity):
         """
         message = {"state": "OFF"}
 
-        if ATTR_TRANSITION in kwargs:
-            message["transition"] = kwargs[ATTR_TRANSITION]
+        self._set_flash_and_transition(message, **kwargs)
 
         mqtt.async_publish(
             self.hass,

tests/components/light/common.py

@@ -111,16 +111,20 @@ async def async_turn_on(
 
 
 @bind_hass
-def turn_off(hass, entity_id=ENTITY_MATCH_ALL, transition=None):
+def turn_off(hass, entity_id=ENTITY_MATCH_ALL, transition=None, flash=None):
     """Turn all or specified light off."""
-    hass.add_job(async_turn_off, hass, entity_id, transition)
+    hass.add_job(async_turn_off, hass, entity_id, transition, flash)
 
 
-async def async_turn_off(hass, entity_id=ENTITY_MATCH_ALL, transition=None):
+async def async_turn_off(hass, entity_id=ENTITY_MATCH_ALL, transition=None, flash=None):
     """Turn all or specified light off."""
     data = {
         key: value
-        for key, value in [(ATTR_ENTITY_ID, entity_id), (ATTR_TRANSITION, transition)]
+        for key, value in [
+            (ATTR_ENTITY_ID, entity_id),
+            (ATTR_TRANSITION, transition),
+            (ATTR_FLASH, flash),
+        ]
         if value is not None
     }
 

tests/components/mqtt/test_light_json.py

@@ -876,6 +876,24 @@ async def test_flash_short_and_long(hass, mqtt_mock):
     state = hass.states.get("light.test")
     assert state.state == STATE_ON
 
+    await common.async_turn_off(hass, "light.test", flash="short")
+
+    mqtt_mock.async_publish.assert_called_once_with(
+        "test_light_rgb/set", JsonValidator('{"state": "OFF", "flash": 5}'), 0, False
+    )
+    mqtt_mock.async_publish.reset_mock()
+    state = hass.states.get("light.test")
+    assert state.state == STATE_OFF
+
+    await common.async_turn_off(hass, "light.test", flash="long")
+
+    mqtt_mock.async_publish.assert_called_once_with(
+        "test_light_rgb/set", JsonValidator('{"state": "OFF", "flash": 15}'), 0, False
+    )
+    mqtt_mock.async_publish.reset_mock()
+    state = hass.states.get("light.test")
+    assert state.state == STATE_OFF
+
 
 async def test_transition(hass, mqtt_mock):
     """Test for transition time being sent when included."""