Cleanup ManualTriggerSensorEntity (#98629)

* Cleanup ManualTriggerSensorEntity

* ConfigType

homeassistant/components/sql/sensor.py

@@ -19,12 +19,7 @@ from homeassistant.components.recorder import (
     SupportedDialect,
     get_instance,
 )
-from homeassistant.components.sensor import (
-    CONF_STATE_CLASS,
-    SensorDeviceClass,
-    SensorEntity,
-    SensorStateClass,
-)
+from homeassistant.components.sensor import CONF_STATE_CLASS
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_DEVICE_CLASS,
@@ -44,7 +39,7 @@ from homeassistant.helpers.template import Template
 from homeassistant.helpers.template_entity import (
     CONF_AVAILABILITY,
     CONF_PICTURE,
-    ManualTriggerEntity,
+    ManualTriggerSensorEntity,
 )
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
@@ -56,6 +51,16 @@ _LOGGER = logging.getLogger(__name__)
 
 _SQL_LAMBDA_CACHE: LRUCache = LRUCache(1000)
 
+TRIGGER_ENTITY_OPTIONS = (
+    CONF_AVAILABILITY,
+    CONF_DEVICE_CLASS,
+    CONF_ICON,
+    CONF_PICTURE,
+    CONF_UNIQUE_ID,
+    CONF_STATE_CLASS,
+    CONF_UNIT_OF_MEASUREMENT,
+)
+
 
 async def async_setup_platform(
     hass: HomeAssistant,
@@ -69,43 +74,29 @@ async def async_setup_platform(
 
     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)
     column_name: str = conf[CONF_COLUMN_NAME]
     unique_id: str | None = conf.get(CONF_UNIQUE_ID)
     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,
-        CONF_UNIQUE_ID: unique_id,
-    }
-    if availability:
-        trigger_entity_config[CONF_AVAILABILITY] = availability
-    if icon:
-        trigger_entity_config[CONF_ICON] = icon
-    if picture:
-        trigger_entity_config[CONF_PICTURE] = picture
+    trigger_entity_config = {CONF_NAME: name}
+    for key in TRIGGER_ENTITY_OPTIONS:
+        if key not in conf:
+            continue
+        trigger_entity_config[key] = conf[key]
 
     await async_setup_sensor(
         hass,
         trigger_entity_config,
         query_str,
         column_name,
-        unit,
         value_template,
         unique_id,
         db_url,
         True,
-        state_class,
         async_add_entities,
     )
 
@@ -118,11 +109,8 @@ async def async_setup_entry(
     db_url: str = resolve_db_url(hass, entry.options.get(CONF_DB_URL))
     name: str = entry.options[CONF_NAME]
     query_str: str = entry.options[CONF_QUERY]
-    unit: str | None = entry.options.get(CONF_UNIT_OF_MEASUREMENT)
     template: str | None = entry.options.get(CONF_VALUE_TEMPLATE)
     column_name: str = entry.options[CONF_COLUMN_NAME]
-    device_class: SensorDeviceClass | None = entry.options.get(CONF_DEVICE_CLASS, None)
-    state_class: SensorStateClass | None = entry.options.get(CONF_STATE_CLASS, None)
 
     value_template: Template | None = None
     if template is not None:
@@ -135,23 +123,21 @@ async def async_setup_entry(
             value_template.hass = hass
 
     name_template = Template(name, hass)
-    trigger_entity_config = {
-        CONF_NAME: name_template,
-        CONF_DEVICE_CLASS: device_class,
-        CONF_UNIQUE_ID: entry.entry_id,
-    }
+    trigger_entity_config = {CONF_NAME: name_template}
+    for key in TRIGGER_ENTITY_OPTIONS:
+        if key not in entry.options:
+            continue
+        trigger_entity_config[key] = entry.options[key]
 
     await async_setup_sensor(
         hass,
         trigger_entity_config,
         query_str,
         column_name,
-        unit,
         value_template,
         entry.entry_id,
         db_url,
         False,
-        state_class,
         async_add_entities,
     )
 
@@ -191,12 +177,10 @@ async def async_setup_sensor(
     trigger_entity_config: ConfigType,
     query_str: str,
     column_name: str,
-    unit: str | None,
     value_template: Template | None,
     unique_id: str | None,
     db_url: str,
     yaml: bool,
-    state_class: SensorStateClass | None,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the SQL sensor."""
@@ -274,10 +258,8 @@ async def async_setup_sensor(
                 sessmaker,
                 query_str,
                 column_name,
-                unit,
                 value_template,
                 yaml,
-                state_class,
                 use_database_executor,
             )
         ],
@@ -317,7 +299,7 @@ def _generate_lambda_stmt(query: str) -> StatementLambdaElement:
     return lambda_stmt(lambda: text, lambda_cache=_SQL_LAMBDA_CACHE)
 
 
-class SQLSensor(ManualTriggerEntity, SensorEntity):
+class SQLSensor(ManualTriggerSensorEntity):
     """Representation of an SQL sensor."""
 
     def __init__(
@@ -326,17 +308,13 @@ class SQLSensor(ManualTriggerEntity, SensorEntity):
         sessmaker: scoped_session,
         query: str,
         column: str,
-        unit: str | None,
         value_template: Template | None,
         yaml: bool,
-        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_native_unit_of_measurement = unit
-        self._attr_state_class = state_class
         self._template = value_template
         self._column_name = column
         self.sessionmaker = sessmaker