From 4e2b00a44369365066868d78f4664b1b93e5b057 Mon Sep 17 00:00:00 2001 From: G Johansson Date: Thu, 20 Jul 2023 11:35:08 +0200 Subject: [PATCH] Refactor SQL with ManualTriggerEntity (#95116) * First go * Finalize sensor * Add tests * Remove not need _attr_name * device_class * _process_manual_data allow Any as value --- homeassistant/components/sql/__init__.py | 7 ++- homeassistant/components/sql/sensor.py | 68 +++++++++++++++-------- homeassistant/helpers/template_entity.py | 2 +- tests/components/sql/__init__.py | 19 +++++++ tests/components/sql/test_sensor.py | 70 +++++++++++++++++++++++- 5 files changed, 140 insertions(+), 26 deletions(-) diff --git a/homeassistant/components/sql/__init__.py b/homeassistant/components/sql/__init__.py index dd5480450e2..316e816fd6f 100644 --- a/homeassistant/components/sql/__init__.py +++ b/homeassistant/components/sql/__init__.py @@ -14,6 +14,7 @@ from homeassistant.components.sensor import ( from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( CONF_DEVICE_CLASS, + CONF_ICON, CONF_NAME, CONF_UNIQUE_ID, CONF_UNIT_OF_MEASUREMENT, @@ -23,6 +24,7 @@ from homeassistant.const import ( from homeassistant.core import HomeAssistant from homeassistant.helpers import discovery import homeassistant.helpers.config_validation as cv +from homeassistant.helpers.template_entity import CONF_AVAILABILITY, CONF_PICTURE from homeassistant.helpers.typing import ConfigType from .const import CONF_COLUMN_NAME, CONF_QUERY, DOMAIN, PLATFORMS @@ -41,7 +43,7 @@ def validate_sql_select(value: str) -> str: QUERY_SCHEMA = vol.Schema( { vol.Required(CONF_COLUMN_NAME): cv.string, - vol.Required(CONF_NAME): cv.string, + vol.Required(CONF_NAME): cv.template, vol.Required(CONF_QUERY): vol.All(cv.string, validate_sql_select), vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string, vol.Optional(CONF_VALUE_TEMPLATE): cv.template, @@ -49,6 +51,9 @@ QUERY_SCHEMA = vol.Schema( vol.Optional(CONF_DB_URL): cv.string, vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA, vol.Optional(CONF_STATE_CLASS): STATE_CLASSES_SCHEMA, + vol.Optional(CONF_AVAILABILITY): cv.template, + vol.Optional(CONF_ICON): cv.template, + vol.Optional(CONF_PICTURE): cv.template, } ) diff --git a/homeassistant/components/sql/sensor.py b/homeassistant/components/sql/sensor.py index 96fc4bc943a..cbdef90f623 100644 --- a/homeassistant/components/sql/sensor.py +++ b/homeassistant/components/sql/sensor.py @@ -4,6 +4,7 @@ from __future__ import annotations from datetime import date import decimal import logging +from typing import Any import sqlalchemy from sqlalchemy import lambda_stmt @@ -27,6 +28,7 @@ from homeassistant.components.sensor import ( from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( CONF_DEVICE_CLASS, + CONF_ICON, CONF_NAME, CONF_UNIQUE_ID, CONF_UNIT_OF_MEASUREMENT, @@ -40,6 +42,11 @@ from homeassistant.helpers.device_registry import DeviceEntryType from homeassistant.helpers.entity import DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.template import Template +from homeassistant.helpers.template_entity import ( + CONF_AVAILABILITY, + CONF_PICTURE, + ManualTriggerEntity, +) from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType from .const import CONF_COLUMN_NAME, CONF_QUERY, DOMAIN @@ -61,7 +68,7 @@ async def async_setup_platform( if (conf := discovery_info) is None: return - name: str = conf[CONF_NAME] + name: Template = conf[CONF_NAME] query_str: str = conf[CONF_QUERY] unit: str | None = conf.get(CONF_UNIT_OF_MEASUREMENT) value_template: Template | None = conf.get(CONF_VALUE_TEMPLATE) @@ -70,13 +77,24 @@ async def async_setup_platform( db_url: str = resolve_db_url(hass, conf.get(CONF_DB_URL)) device_class: SensorDeviceClass | None = conf.get(CONF_DEVICE_CLASS) state_class: SensorStateClass | None = conf.get(CONF_STATE_CLASS) + availability: Template | None = conf.get(CONF_AVAILABILITY) + icon: Template | None = conf.get(CONF_ICON) + picture: Template | None = conf.get(CONF_PICTURE) if value_template is not None: value_template.hass = hass + trigger_entity_config = {CONF_NAME: name, CONF_DEVICE_CLASS: device_class} + if availability: + trigger_entity_config[CONF_AVAILABILITY] = availability + if icon: + trigger_entity_config[CONF_ICON] = icon + if picture: + trigger_entity_config[CONF_PICTURE] = picture + await async_setup_sensor( hass, - name, + trigger_entity_config, query_str, column_name, unit, @@ -84,7 +102,6 @@ async def async_setup_platform( unique_id, db_url, True, - device_class, state_class, async_add_entities, ) @@ -114,9 +131,12 @@ async def async_setup_entry( if value_template is not None: value_template.hass = hass + name_template = Template(name, hass) + trigger_entity_config = {CONF_NAME: name_template, CONF_DEVICE_CLASS: device_class} + await async_setup_sensor( hass, - name, + trigger_entity_config, query_str, column_name, unit, @@ -124,7 +144,6 @@ async def async_setup_entry( entry.entry_id, db_url, False, - device_class, state_class, async_add_entities, ) @@ -162,7 +181,7 @@ def _async_get_or_init_domain_data(hass: HomeAssistant) -> SQLData: async def async_setup_sensor( hass: HomeAssistant, - name: str, + trigger_entity_config: ConfigType, query_str: str, column_name: str, unit: str | None, @@ -170,7 +189,6 @@ async def async_setup_sensor( unique_id: str | None, db_url: str, yaml: bool, - device_class: SensorDeviceClass | None, state_class: SensorStateClass | None, async_add_entities: AddEntitiesCallback, ) -> None: @@ -245,7 +263,7 @@ async def async_setup_sensor( async_add_entities( [ SQLSensor( - name, + trigger_entity_config, sessmaker, query_str, column_name, @@ -253,12 +271,10 @@ async def async_setup_sensor( value_template, unique_id, yaml, - device_class, state_class, use_database_executor, ) ], - True, ) @@ -295,15 +311,12 @@ def _generate_lambda_stmt(query: str) -> StatementLambdaElement: return lambda_stmt(lambda: text, lambda_cache=_SQL_LAMBDA_CACHE) -class SQLSensor(SensorEntity): +class SQLSensor(ManualTriggerEntity, SensorEntity): """Representation of an SQL sensor.""" - _attr_icon = "mdi:database-search" - _attr_has_entity_name = True - def __init__( self, - name: str, + trigger_entity_config: ConfigType, sessmaker: scoped_session, query: str, column: str, @@ -311,15 +324,13 @@ class SQLSensor(SensorEntity): value_template: Template | None, unique_id: str | None, yaml: bool, - device_class: SensorDeviceClass | None, state_class: SensorStateClass | None, use_database_executor: bool, ) -> None: """Initialize the SQL sensor.""" + super().__init__(self.hass, trigger_entity_config) self._query = query - self._attr_name = name if yaml else None self._attr_native_unit_of_measurement = unit - self._attr_device_class = device_class self._attr_state_class = state_class self._template = value_template self._column_name = column @@ -328,22 +339,34 @@ class SQLSensor(SensorEntity): self._attr_unique_id = unique_id self._use_database_executor = use_database_executor self._lambda_stmt = _generate_lambda_stmt(query) + self._attr_has_entity_name = not yaml if not yaml and unique_id: self._attr_device_info = DeviceInfo( entry_type=DeviceEntryType.SERVICE, identifiers={(DOMAIN, unique_id)}, manufacturer="SQL", - name=name, + name=trigger_entity_config[CONF_NAME].template, ) + async def async_added_to_hass(self) -> None: + """Call when entity about to be added to hass.""" + await super().async_added_to_hass() + await self.async_update() + + @property + def extra_state_attributes(self) -> dict[str, Any] | None: + """Return extra attributes.""" + return dict(self._attr_extra_state_attributes) + async def async_update(self) -> None: """Retrieve sensor data from the query using the right executor.""" if self._use_database_executor: - await get_instance(self.hass).async_add_executor_job(self._update) + data = await get_instance(self.hass).async_add_executor_job(self._update) else: - await self.hass.async_add_executor_job(self._update) + data = await self.hass.async_add_executor_job(self._update) + self._process_manual_data(data) - def _update(self) -> None: + def _update(self) -> Any: """Retrieve sensor data from the query.""" data = None self._attr_extra_state_attributes = {} @@ -384,3 +407,4 @@ class SQLSensor(SensorEntity): _LOGGER.warning("%s returned no results", self._query) sess.close() + return data diff --git a/homeassistant/helpers/template_entity.py b/homeassistant/helpers/template_entity.py index 42d578555ab..fcd98a77831 100644 --- a/homeassistant/helpers/template_entity.py +++ b/homeassistant/helpers/template_entity.py @@ -624,7 +624,7 @@ class ManualTriggerEntity(TriggerBaseEntity): TriggerBaseEntity.__init__(self, hass, config) @callback - def _process_manual_data(self, value: str | None = None) -> None: + def _process_manual_data(self, value: Any | None = None) -> None: """Process new data manually. Implementing class should call this last in update method to render templates. diff --git a/tests/components/sql/__init__.py b/tests/components/sql/__init__.py index a1417cd38df..53356a85c4e 100644 --- a/tests/components/sql/__init__.py +++ b/tests/components/sql/__init__.py @@ -13,12 +13,14 @@ from homeassistant.components.sql.const import CONF_COLUMN_NAME, CONF_QUERY, DOM from homeassistant.config_entries import SOURCE_USER from homeassistant.const import ( CONF_DEVICE_CLASS, + CONF_ICON, CONF_NAME, CONF_UNIQUE_ID, CONF_UNIT_OF_MEASUREMENT, CONF_VALUE_TEMPLATE, ) from homeassistant.core import HomeAssistant +from homeassistant.helpers.template_entity import CONF_AVAILABILITY, CONF_PICTURE from tests.common import MockConfigEntry @@ -148,6 +150,23 @@ YAML_CONFIG_NO_DB = { } } +YAML_CONFIG_ALL_TEMPLATES = { + "sql": { + CONF_DB_URL: "sqlite://", + CONF_NAME: "Get values with template", + CONF_QUERY: "SELECT 5 as output", + CONF_COLUMN_NAME: "output", + CONF_UNIT_OF_MEASUREMENT: "MiB/s", + CONF_UNIQUE_ID: "unique_id_123456", + CONF_VALUE_TEMPLATE: "{{ value }}", + CONF_ICON: '{% if states("sensor.input1")=="on" %} mdi:on {% else %} mdi:off {% endif %}', + CONF_PICTURE: '{% if states("sensor.input1")=="on" %} /local/picture1.jpg {% else %} /local/picture2.jpg {% endif %}', + CONF_AVAILABILITY: '{{ states("sensor.input2")=="on" }}', + CONF_DEVICE_CLASS: SensorDeviceClass.DATA_RATE, + CONF_STATE_CLASS: SensorStateClass.MEASUREMENT, + } +} + async def init_integration( hass: HomeAssistant, diff --git a/tests/components/sql/test_sensor.py b/tests/components/sql/test_sensor.py index 0fe0e881c95..3d0e2768ade 100644 --- a/tests/components/sql/test_sensor.py +++ b/tests/components/sql/test_sensor.py @@ -13,7 +13,12 @@ from homeassistant.components.recorder import Recorder from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass from homeassistant.components.sql.const import CONF_QUERY, DOMAIN from homeassistant.config_entries import SOURCE_USER -from homeassistant.const import CONF_UNIQUE_ID, STATE_UNKNOWN +from homeassistant.const import ( + CONF_ICON, + CONF_UNIQUE_ID, + STATE_UNAVAILABLE, + STATE_UNKNOWN, +) from homeassistant.core import HomeAssistant from homeassistant.helpers import issue_registry as ir from homeassistant.setup import async_setup_component @@ -21,6 +26,7 @@ from homeassistant.util import dt as dt_util from . import ( YAML_CONFIG, + YAML_CONFIG_ALL_TEMPLATES, YAML_CONFIG_BINARY, YAML_CONFIG_FULL_TABLE_SCAN, YAML_CONFIG_FULL_TABLE_SCAN_NO_UNIQUE_ID, @@ -32,13 +38,14 @@ from . import ( from tests.common import MockConfigEntry, async_fire_time_changed -async def test_query(recorder_mock: Recorder, hass: HomeAssistant) -> None: +async def test_query_basic(recorder_mock: Recorder, hass: HomeAssistant) -> None: """Test the SQL sensor.""" config = { "db_url": "sqlite://", "query": "SELECT 5 as value", "column": "value", "name": "Select value SQL query", + "unique_id": "very_unique_id", } await init_integration(hass, config) @@ -235,6 +242,65 @@ async def test_query_from_yaml(recorder_mock: Recorder, hass: HomeAssistant) -> assert state.state == "5" +async def test_templates_with_yaml( + recorder_mock: Recorder, hass: HomeAssistant +) -> None: + """Test the SQL sensor from yaml config with templates.""" + + hass.states.async_set("sensor.input1", "on") + hass.states.async_set("sensor.input2", "on") + await hass.async_block_till_done() + + assert await async_setup_component(hass, DOMAIN, YAML_CONFIG_ALL_TEMPLATES) + await hass.async_block_till_done() + + state = hass.states.get("sensor.get_values_with_template") + assert state.state == "5" + assert state.attributes[CONF_ICON] == "mdi:on" + assert state.attributes["entity_picture"] == "/local/picture1.jpg" + + hass.states.async_set("sensor.input1", "off") + await hass.async_block_till_done() + + async_fire_time_changed( + hass, + dt_util.utcnow() + timedelta(minutes=1), + ) + await hass.async_block_till_done() + + state = hass.states.get("sensor.get_values_with_template") + assert state.state == "5" + assert state.attributes[CONF_ICON] == "mdi:off" + assert state.attributes["entity_picture"] == "/local/picture2.jpg" + + hass.states.async_set("sensor.input2", "off") + await hass.async_block_till_done() + + async_fire_time_changed( + hass, + dt_util.utcnow() + timedelta(minutes=2), + ) + await hass.async_block_till_done() + + state = hass.states.get("sensor.get_values_with_template") + assert state.state == STATE_UNAVAILABLE + + hass.states.async_set("sensor.input1", "on") + hass.states.async_set("sensor.input2", "on") + await hass.async_block_till_done() + + async_fire_time_changed( + hass, + dt_util.utcnow() + timedelta(minutes=3), + ) + await hass.async_block_till_done() + + state = hass.states.get("sensor.get_values_with_template") + assert state.state == "5" + assert state.attributes[CONF_ICON] == "mdi:on" + assert state.attributes["entity_picture"] == "/local/picture1.jpg" + + async def test_config_from_old_yaml( recorder_mock: Recorder, hass: HomeAssistant ) -> None: