Update websocket api to use async_track_template_result (#39057)

This commit is contained in:
J. Nick Koston 2020-08-21 07:04:29 -05:00 committed by GitHub
parent 7878d97588
commit f560256546
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 51 additions and 23 deletions

View File

@ -1,5 +1,6 @@
"""Commands part of Websocket API."""
import asyncio
import logging
import voluptuous as vol
@ -7,14 +8,21 @@ from homeassistant.auth.permissions.const import CAT_ENTITIES, POLICY_READ
from homeassistant.components.websocket_api.const import ERR_NOT_FOUND
from homeassistant.const import EVENT_STATE_CHANGED, EVENT_TIME_CHANGED, MATCH_ALL
from homeassistant.core import DOMAIN as HASS_DOMAIN, callback
from homeassistant.exceptions import HomeAssistantError, ServiceNotFound, Unauthorized
from homeassistant.exceptions import (
HomeAssistantError,
ServiceNotFound,
TemplateError,
Unauthorized,
)
from homeassistant.helpers import config_validation as cv, entity
from homeassistant.helpers.event import async_track_state_change_event
from homeassistant.helpers.event import async_track_template_result
from homeassistant.helpers.service import async_get_all_descriptions
from homeassistant.loader import IntegrationNotFound, async_get_integration
from . import const, decorators, messages
_LOGGER = logging.getLogger(__name__)
# mypy: allow-untyped-calls, allow-untyped-defs
@ -244,27 +252,26 @@ def handle_render_template(hass, connection, msg):
variables = msg.get("variables")
entity_ids = msg.get("entity_ids")
if entity_ids is None:
entity_ids = template.extract_entities(variables)
@callback
def state_listener(*_):
connection.send_message(
messages.event_message(
msg["id"], {"result": template.async_render(variables)}
def _template_listener(event, template, last_result, result):
if isinstance(result, TemplateError):
_LOGGER.error(
"TemplateError('%s') " "while processing template '%s'",
result,
template,
)
)
if entity_ids and entity_ids != MATCH_ALL:
connection.subscriptions[msg["id"]] = async_track_state_change_event(
hass, entity_ids, state_listener
)
else:
connection.subscriptions[msg["id"]] = lambda: None
result = None
connection.send_message(messages.event_message(msg["id"], {"result": result}))
info = async_track_template_result(hass, template, _template_listener, variables)
connection.subscriptions[msg["id"]] = info.async_remove
connection.send_result(msg["id"])
state_listener()
hass.loop.call_soon_threadsafe(info.async_refresh)
@callback

View File

@ -430,19 +430,17 @@ async def test_render_template_renders_template(
assert event == {"result": "State is: off"}
async def test_render_template_with_manual_entity_ids(
async def test_render_template_manual_entity_ids_no_longer_needed(
hass, websocket_client, hass_admin_user
):
"""Test that updates to specified entity ids cause a template rerender."""
hass.states.async_set("light.test", "on")
hass.states.async_set("light.test2", "on")
await websocket_client.send_json(
{
"id": 5,
"type": "render_template",
"template": "State is: {{ states('light.test') }}",
"entity_ids": ["light.test2"],
}
)
@ -457,12 +455,35 @@ async def test_render_template_with_manual_entity_ids(
event = msg["event"]
assert event == {"result": "State is: on"}
hass.states.async_set("light.test2", "off")
hass.states.async_set("light.test", "off")
msg = await websocket_client.receive_json()
assert msg["id"] == 5
assert msg["type"] == "event"
event = msg["event"]
assert event == {"result": "State is: on"}
assert event == {"result": "State is: off"}
async def test_render_template_with_error(
hass, websocket_client, hass_admin_user, caplog
):
"""Test a template with an error."""
await websocket_client.send_json(
{"id": 5, "type": "render_template", "template": "{{ my_unknown_var() + 1 }}"}
)
msg = await websocket_client.receive_json()
assert msg["id"] == 5
assert msg["type"] == const.TYPE_RESULT
assert msg["success"]
msg = await websocket_client.receive_json()
assert msg["id"] == 5
assert msg["type"] == "event"
event = msg["event"]
assert event == {"result": None}
assert "my_unknown_var" in caplog.text
assert "TemplateError" in caplog.text
async def test_render_template_returns_with_match_all(