Add availability_template to Template Switch platform

homeassistant/components/template/switch.py

@@ -19,12 +19,14 @@
     ATTR_ENTITY_ID,
     CONF_SWITCHES,
     EVENT_HOMEASSISTANT_START,
+    MATCH_ALL,
 )
 from homeassistant.exceptions import TemplateError
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import async_generate_entity_id
 from homeassistant.helpers.event import async_track_state_change
 from homeassistant.helpers.script import Script
+from .const import CONF_AVAILABILITY_TEMPLATE
 
 _LOGGER = logging.getLogger(__name__)
 _VALID_STATES = [STATE_ON, STATE_OFF, "true", "false"]
@@ -37,6 +39,7 @@
         vol.Required(CONF_VALUE_TEMPLATE): cv.template,
         vol.Optional(CONF_ICON_TEMPLATE): cv.template,
         vol.Optional(CONF_ENTITY_PICTURE_TEMPLATE): cv.template,
+        vol.Optional(CONF_AVAILABILITY_TEMPLATE): cv.template,
         vol.Required(ON_ACTION): cv.SCRIPT_SCHEMA,
         vol.Required(OFF_ACTION): cv.SCRIPT_SCHEMA,
         vol.Optional(ATTR_FRIENDLY_NAME): cv.string,
@@ -58,19 +61,47 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
         state_template = device_config[CONF_VALUE_TEMPLATE]
         icon_template = device_config.get(CONF_ICON_TEMPLATE)
         entity_picture_template = device_config.get(CONF_ENTITY_PICTURE_TEMPLATE)
+        availability_template = device_config.get(CONF_AVAILABILITY_TEMPLATE)
         on_action = device_config[ON_ACTION]
         off_action = device_config[OFF_ACTION]
-        entity_ids = (
-            device_config.get(ATTR_ENTITY_ID) or state_template.extract_entities()
-        )
-
-        state_template.hass = hass
-
-        if icon_template is not None:
-            icon_template.hass = hass
-
-        if entity_picture_template is not None:
-            entity_picture_template.hass = hass
+        manual_entity_ids = device_config.get(ATTR_ENTITY_ID)
+        entity_ids = set()
+
+        templates = {
+            CONF_VALUE_TEMPLATE: state_template,
+            CONF_ICON_TEMPLATE: icon_template,
+            CONF_ENTITY_PICTURE_TEMPLATE: entity_picture_template,
+            CONF_AVAILABILITY_TEMPLATE: availability_template,
+        }
+        invalid_templates = []
+
+        for template_name, template in templates.items():
+            if template is not None:
+                template.hass = hass
+
+                if manual_entity_ids is not None:
+                    continue
+
+                template_entity_ids = template.extract_entities()
+                if template_entity_ids == MATCH_ALL:
+                    invalid_templates.append(template_name.replace("_template", ""))
+                    entity_ids = MATCH_ALL
+                elif entity_ids != MATCH_ALL:
+                    entity_ids |= set(template_entity_ids)
+        if invalid_templates:
+            _LOGGER.warning(
+                "Template sensor %s has no entity ids configured to track nor"
+                " were we able to extract the entities to track from the %s "
+                "template(s). This entity will only be able to be updated "
+                "manually.",
+                device,
+                ", ".join(invalid_templates),
+            )
+        else:
+            if manual_entity_ids is None:
+                entity_ids = list(entity_ids)
+            else:
+                entity_ids = manual_entity_ids
 
         switches.append(
             SwitchTemplate(
@@ -80,6 +111,7 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
                 state_template,
                 icon_template,
                 entity_picture_template,
+                availability_template,
                 on_action,
                 off_action,
                 entity_ids,
@@ -104,6 +136,7 @@ def __init__(
         state_template,
         icon_template,
         entity_picture_template,
+        availability_template,
         on_action,
         off_action,
         entity_ids,
@@ -120,9 +153,11 @@ def __init__(
         self._state = False
         self._icon_template = icon_template
         self._entity_picture_template = entity_picture_template
+        self._availability_template = availability_template
         self._icon = None
         self._entity_picture = None
         self._entities = entity_ids
+        self._available = True
 
     async def async_added_to_hass(self):
         """Register callbacks."""
@@ -160,11 +195,6 @@ def should_poll(self):
         """Return the polling state."""
         return False
 
-    @property
-    def available(self):
-        """If switch is available."""
-        return self._state is not None
-
     @property
     def icon(self):
         """Return the icon to use in the frontend, if any."""
@@ -175,6 +205,11 @@ def entity_picture(self):
         """Return the entity_picture to use in the frontend, if any."""
         return self._entity_picture
 
+    @property
+    def available(self) -> bool:
+        """Return if the device is available."""
+        return self._available
+
     async def async_turn_on(self, **kwargs):
         """Fire the on action."""
         await self._on_script.async_run(context=self._context)
@@ -205,12 +240,16 @@ async def async_update(self):
         for property_name, template in (
             ("_icon", self._icon_template),
             ("_entity_picture", self._entity_picture_template),
+            ("_available", self._availability_template),
         ):
             if template is None:
                 continue
 
             try:
-                setattr(self, property_name, template.async_render())
+                value = template.async_render()
+                if property_name == "_available":
+                    value = value.lower() == "true"
+                setattr(self, property_name, value)
             except TemplateError as ex:
                 friendly_property_name = property_name[1:].replace("_", " ")
                 if ex.args and ex.args[0].startswith(

tests/components/template/test_switch.py

@@ -1,7 +1,7 @@
 """The tests for the  Template switch platform."""
 from homeassistant.core import callback
 from homeassistant import setup
-from homeassistant.const import STATE_ON, STATE_OFF
+from homeassistant.const import STATE_ON, STATE_OFF, STATE_UNAVAILABLE
 
 from tests.common import get_test_home_assistant, assert_setup_component
 from tests.components.switch import common
@@ -474,3 +474,76 @@ def test_off_action(self):
         self.hass.block_till_done()
 
         assert len(self.calls) == 1
+
+
+async def test_available_template_with_entities(hass):
+    """Test availability templates with values from other entities."""
+    await setup.async_setup_component(
+        hass,
+        "switch",
+        {
+            "switch": {
+                "platform": "template",
+                "switches": {
+                    "test_template_switch": {
+                        "value_template": "{{ 1 == 1 }}",
+                        "turn_on": {
+                            "service": "switch.turn_on",
+                            "entity_id": "switch.test_state",
+                        },
+                        "turn_off": {
+                            "service": "switch.turn_off",
+                            "entity_id": "switch.test_state",
+                        },
+                        "availability_template": "{{ is_state('availability_state.state', 'on') }}",
+                    }
+                },
+            }
+        },
+    )
+
+    await hass.async_start()
+    await hass.async_block_till_done()
+
+    hass.states.async_set("availability_state.state", STATE_ON)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("switch.test_template_switch").state != STATE_UNAVAILABLE
+
+    hass.states.async_set("availability_state.state", STATE_OFF)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("switch.test_template_switch").state == STATE_UNAVAILABLE
+
+
+async def test_invalid_availability_template_keeps_component_available(hass, caplog):
+    """Test that an invalid availability keeps the device available."""
+    await setup.async_setup_component(
+        hass,
+        "switch",
+        {
+            "switch": {
+                "platform": "template",
+                "switches": {
+                    "test_template_switch": {
+                        "value_template": "{{ true }}",
+                        "turn_on": {
+                            "service": "switch.turn_on",
+                            "entity_id": "switch.test_state",
+                        },
+                        "turn_off": {
+                            "service": "switch.turn_off",
+                            "entity_id": "switch.test_state",
+                        },
+                        "availability_template": "{{ x - 12 }}",
+                    }
+                },
+            }
+        },
+    )
+
+    await hass.async_start()
+    await hass.async_block_till_done()
+
+    assert hass.states.get("switch.test_template_switch").state != STATE_UNAVAILABLE
+    assert ("UndefinedError: 'x' is undefined") in caplog.text