diff --git a/homeassistant/components/sql/__init__.py b/homeassistant/components/sql/__init__.py index bba49c415f8..c0ec2dfab7f 100644 --- a/homeassistant/components/sql/__init__.py +++ b/homeassistant/components/sql/__init__.py @@ -4,8 +4,14 @@ from __future__ import annotations import voluptuous as vol from homeassistant.components.recorder import CONF_DB_URL +from homeassistant.components.sensor import ( + CONF_STATE_CLASS, + DEVICE_CLASSES_SCHEMA, + STATE_CLASSES_SCHEMA, +) from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( + CONF_DEVICE_CLASS, CONF_NAME, CONF_UNIQUE_ID, CONF_UNIT_OF_MEASUREMENT, @@ -36,6 +42,8 @@ QUERY_SCHEMA = vol.Schema( vol.Optional(CONF_VALUE_TEMPLATE): cv.template, vol.Optional(CONF_UNIQUE_ID): cv.string, vol.Optional(CONF_DB_URL): cv.string, + vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA, + vol.Optional(CONF_STATE_CLASS): STATE_CLASSES_SCHEMA, } ) diff --git a/homeassistant/components/sql/sensor.py b/homeassistant/components/sql/sensor.py index 5d51087a9dd..27cf798db38 100644 --- a/homeassistant/components/sql/sensor.py +++ b/homeassistant/components/sql/sensor.py @@ -11,9 +11,15 @@ from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session, scoped_session, sessionmaker from homeassistant.components.recorder import CONF_DB_URL, DEFAULT_DB_FILE, DEFAULT_URL -from homeassistant.components.sensor import SensorEntity +from homeassistant.components.sensor import ( + CONF_STATE_CLASS, + SensorDeviceClass, + SensorEntity, + SensorStateClass, +) from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( + CONF_DEVICE_CLASS, CONF_NAME, CONF_UNIQUE_ID, CONF_UNIT_OF_MEASUREMENT, @@ -54,6 +60,8 @@ async def async_setup_platform( column_name: str = conf[CONF_COLUMN_NAME] unique_id: str | None = conf.get(CONF_UNIQUE_ID) db_url: str | None = conf.get(CONF_DB_URL) + device_class: SensorDeviceClass | None = conf.get(CONF_DEVICE_CLASS) + state_class: SensorStateClass | None = conf.get(CONF_STATE_CLASS) if value_template is not None: value_template.hass = hass @@ -68,6 +76,8 @@ async def async_setup_platform( unique_id, db_url, True, + device_class, + state_class, async_add_entities, ) @@ -104,6 +114,8 @@ async def async_setup_entry( entry.entry_id, db_url, False, + None, + None, async_add_entities, ) @@ -118,6 +130,8 @@ async def async_setup_sensor( unique_id: str | None, db_url: str | None, yaml: bool, + device_class: SensorDeviceClass | None, + state_class: SensorStateClass | None, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the SQL sensor.""" @@ -163,6 +177,8 @@ async def async_setup_sensor( value_template, unique_id, yaml, + device_class, + state_class, ) ], True, @@ -185,11 +201,15 @@ class SQLSensor(SensorEntity): value_template: Template | None, unique_id: str | None, yaml: bool, + device_class: SensorDeviceClass | None, + state_class: SensorStateClass | None, ) -> None: """Initialize the SQL sensor.""" 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 self.sessionmaker = sessmaker diff --git a/tests/components/sql/__init__.py b/tests/components/sql/__init__.py index 1c096dfef6a..f6cfba01e35 100644 --- a/tests/components/sql/__init__.py +++ b/tests/components/sql/__init__.py @@ -4,9 +4,15 @@ from __future__ import annotations from typing import Any from homeassistant.components.recorder import CONF_DB_URL +from homeassistant.components.sensor import ( + CONF_STATE_CLASS, + SensorDeviceClass, + SensorStateClass, +) from homeassistant.components.sql.const import CONF_COLUMN_NAME, CONF_QUERY, DOMAIN from homeassistant.config_entries import SOURCE_USER from homeassistant.const import ( + CONF_DEVICE_CLASS, CONF_NAME, CONF_UNIQUE_ID, CONF_UNIT_OF_MEASUREMENT, @@ -56,6 +62,8 @@ YAML_CONFIG = { CONF_UNIT_OF_MEASUREMENT: "MiB", CONF_UNIQUE_ID: "unique_id_12345", CONF_VALUE_TEMPLATE: "{{ value }}", + CONF_DEVICE_CLASS: SensorDeviceClass.DATA_RATE, + CONF_STATE_CLASS: SensorStateClass.MEASUREMENT, } } diff --git a/tests/components/sql/test_sensor.py b/tests/components/sql/test_sensor.py index d7a12795ed0..bc3143347b5 100644 --- a/tests/components/sql/test_sensor.py +++ b/tests/components/sql/test_sensor.py @@ -9,6 +9,7 @@ from sqlalchemy import text as sql_text from sqlalchemy.exc import SQLAlchemyError from homeassistant.components.recorder import Recorder +from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass from homeassistant.components.sql.const import DOMAIN from homeassistant.config_entries import SOURCE_USER from homeassistant.const import STATE_UNKNOWN @@ -300,3 +301,19 @@ async def test_invalid_url_setup_from_yaml( assert pattern not in caplog.text for pattern in expected_patterns: assert pattern in caplog.text + + +async def test_attributes_from_yaml_setup( + recorder_mock: Recorder, hass: HomeAssistant +) -> None: + """Test attributes from yaml config.""" + + assert await async_setup_component(hass, DOMAIN, YAML_CONFIG) + await hass.async_block_till_done() + + state = hass.states.get("sensor.get_value") + + assert state.state == "5" + assert state.attributes["device_class"] == SensorDeviceClass.DATA_RATE + assert state.attributes["state_class"] == SensorStateClass.MEASUREMENT + assert state.attributes["unit_of_measurement"] == "MiB"