Update universal media_player to use async_track_template_result (#39054)

* Update universal media_player to use async_track_template_result

* Review comments and add missing test cover

homeassistant/components/universal/media_player.py

@@ -68,7 +68,8 @@ from homeassistant.const import (
     STATE_ON,
     STATE_UNAVAILABLE,
 )
-from homeassistant.core import callback
+from homeassistant.core import EVENT_HOMEASSISTANT_START, callback
+from homeassistant.exceptions import TemplateError
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.service import async_call_from_config
 
@@ -132,27 +133,45 @@ class UniversalMediaPlayer(MediaPlayerEntity):
                 attr.append(None)
             self._attrs[key] = attr
         self._child_state = None
+        self._state_template_result = None
         self._state_template = state_template
-        if state_template is not None:
-            self._state_template.hass = hass
 
     async def async_added_to_hass(self):
         """Subscribe to children and template state changes."""
 
         @callback
-        def async_on_dependency_update(*_):
+        def _async_on_dependency_update(*_):
             """Update ha state when dependencies update."""
             self.async_schedule_update_ha_state(True)
 
+        @callback
+        def _async_on_template_update(event, template, last_result, result):
+            """Update ha state when dependencies update."""
+            if isinstance(result, TemplateError):
+                self._state_template_result = None
+            else:
+                self._state_template_result = result
+            self.async_schedule_update_ha_state(True)
+
+        if self._state_template is not None:
+            result = self.hass.helpers.event.async_track_template_result(
+                self._state_template, _async_on_template_update
+            )
+
+            self.hass.bus.async_listen_once(
+                EVENT_HOMEASSISTANT_START, callback(lambda _: result.async_refresh())
+            )
+
+            self.async_on_remove(result.async_remove)
+
         depend = copy(self._children)
         for entity in self._attrs.values():
             depend.append(entity[0])
-        if self._state_template is not None:
-            for entity in self._state_template.extract_entities():
-                depend.append(entity)
 
-        self.hass.helpers.event.async_track_state_change_event(
+        self.async_on_remove(
-            list(set(depend)), async_on_dependency_update
+            self.hass.helpers.event.async_track_state_change_event(
+                list(set(depend)), _async_on_dependency_update
+            )
         )
 
     def _entity_lkp(self, entity_id, state_attr=None):
@@ -217,7 +236,7 @@ class UniversalMediaPlayer(MediaPlayerEntity):
     def master_state(self):
         """Return the master state for entity or None."""
         if self._state_template is not None:
-            return self._state_template.async_render()
+            return self._state_template_result
         if CONF_STATE in self._attrs:
             master_state = self._entity_lkp(
                 self._attrs[CONF_STATE][0], self._attrs[CONF_STATE][1]

tests/components/universal/test_media_player.py

@@ -10,7 +10,14 @@ import homeassistant.components.input_select as input_select
 import homeassistant.components.media_player as media_player
 import homeassistant.components.switch as switch
 import homeassistant.components.universal.media_player as universal
-from homeassistant.const import STATE_OFF, STATE_ON, STATE_PAUSED, STATE_PLAYING
+from homeassistant.const import (
+    STATE_OFF,
+    STATE_ON,
+    STATE_PAUSED,
+    STATE_PLAYING,
+    STATE_UNKNOWN,
+)
+from homeassistant.setup import async_setup_component
 
 from tests.common import get_test_home_assistant, mock_service
 
@@ -337,23 +344,6 @@ class TestMediaPlayer(unittest.TestCase):
         self.hass.states.set(self.mock_state_switch_id, STATE_ON)
         assert STATE_ON == ump.master_state
 
-    def test_master_state_with_template(self):
-        """Test the state_template option."""
-        config = copy(self.config_children_and_attr)
-        self.hass.states.set("input_boolean.test", STATE_OFF)
-        templ = (
-            '{% if states.input_boolean.test.state == "off" %}on'
-            "{% else %}{{ states.media_player.mock1.state }}{% endif %}"
-        )
-        config["state_template"] = templ
-        config = validate_config(config)
-
-        ump = universal.UniversalMediaPlayer(self.hass, **config)
-
-        assert STATE_ON == ump.master_state
-        self.hass.states.set("input_boolean.test", STATE_ON)
-        assert STATE_OFF == ump.master_state
-
     def test_master_state_with_bad_attrs(self):
         """Test master state property."""
         config = copy(self.config_children_and_attr)
@@ -735,3 +725,90 @@ class TestMediaPlayer(unittest.TestCase):
 
         asyncio.run_coroutine_threadsafe(ump.async_turn_off(), self.hass.loop).result()
         assert 1 == len(service)
+
+
+async def test_state_template(hass):
+    """Test with a simple valid state template."""
+    hass.states.async_set("sensor.test_sensor", STATE_ON)
+
+    await async_setup_component(
+        hass,
+        "media_player",
+        {
+            "media_player": {
+                "platform": "universal",
+                "name": "tv",
+                "state_template": "{{ states.sensor.test_sensor.state }}",
+            }
+        },
+    )
+    await hass.async_block_till_done()
+    assert len(hass.states.async_all()) == 2
+    await hass.async_start()
+
+    await hass.async_block_till_done()
+    assert hass.states.get("media_player.tv").state == STATE_ON
+    hass.states.async_set("sensor.test_sensor", STATE_OFF)
+    await hass.async_block_till_done()
+    assert hass.states.get("media_player.tv").state == STATE_OFF
+
+
+async def test_invalid_state_template(hass):
+    """Test invalid state template sets state to None."""
+    hass.states.async_set("sensor.test_sensor", "on")
+
+    await async_setup_component(
+        hass,
+        "media_player",
+        {
+            "media_player": {
+                "platform": "universal",
+                "name": "tv",
+                "state_template": "{{ states.sensor.test_sensor.state + x }}",
+            }
+        },
+    )
+    await hass.async_block_till_done()
+    assert len(hass.states.async_all()) == 2
+    await hass.async_start()
+
+    await hass.async_block_till_done()
+    assert hass.states.get("media_player.tv").state == STATE_UNKNOWN
+    hass.states.async_set("sensor.test_sensor", "off")
+    await hass.async_block_till_done()
+    assert hass.states.get("media_player.tv").state == STATE_UNKNOWN
+
+
+async def test_master_state_with_template(hass):
+    """Test the state_template option."""
+    hass.states.async_set("input_boolean.test", STATE_OFF)
+    hass.states.async_set("media_player.mock1", STATE_OFF)
+
+    templ = (
+        '{% if states.input_boolean.test.state == "off" %}on'
+        "{% else %}{{ states.media_player.mock1.state }}{% endif %}"
+    )
+
+    await async_setup_component(
+        hass,
+        "media_player",
+        {
+            "media_player": {
+                "platform": "universal",
+                "name": "tv",
+                "state_template": templ,
+            }
+        },
+    )
+
+    await hass.async_block_till_done()
+    assert len(hass.states.async_all()) == 3
+    await hass.async_start()
+
+    await hass.async_block_till_done()
+    hass.states.get("media_player.tv").state == STATE_ON
+
+    hass.states.async_set("input_boolean.test", STATE_ON)
+    await hass.async_block_till_done()
+
+    hass.states.get("media_player.tv").state == STATE_OFF