Exclude unimportant attributes from being recorded in the database (#132

)

custom_components/average/const.py

@@ -10,7 +10,7 @@
 # Base component constants
 NAME: Final = "Average Sensor"
 DOMAIN: Final = "average"
-VERSION: Final = "2.2.3"
+VERSION: Final = "2.2.4-alpha"
 ISSUE_URL: Final = "https://github.com/Limych/ha-average/issues"
 
 STARTUP_MESSAGE: Final = f"""

custom_components/average/manifest.json

@@ -1,7 +1,7 @@
 {
   "domain": "average",
   "name": "Average Sensor",
-  "version": "2.2.3",
+  "version": "2.2.4-alpha",
   "documentation": "https://github.com/Limych/ha-average",
   "issue_tracker": "https://github.com/Limych/ha-average/issues",
   "dependencies": [],

custom_components/average/recorder.py

@@ -0,0 +1,14 @@
+"""Integration platform for recorder."""
+from __future__ import annotations
+
+from homeassistant.core import HomeAssistant, callback
+
+from .const import ATTR_SOURCES, ATTR_TO_PROPERTY
+
+
+@callback
+def exclude_attributes(hass: HomeAssistant) -> set[str]:
+    """Exclude unimportant attributes from being recorded in the database."""
+    attributes_to_exclude = set(ATTR_TO_PROPERTY)
+    attributes_to_exclude.discard(ATTR_SOURCES)
+    return attributes_to_exclude

custom_components/average/sensor.py

@@ -1,4 +1,4 @@
-#  Copyright (c) 2019-2021, Andrey "Limych" Khrolenok <[email protected]>
+#  Copyright (c) 2019-2022, Andrey "Limych" Khrolenok <[email protected]>
 #  Creative Commons BY-NC-SA 4.0 International Public License
 #  (see LICENSE.md or https://creativecommons.org/licenses/by-nc-sa/4.0/)
 
@@ -8,6 +8,8 @@
 For more details about this sensor, please refer to the documentation at
 https://github.com/Limych/ha-average/
 """
+from __future__ import annotations
+
 import datetime
 import logging
 import math
@@ -19,6 +21,7 @@
 from _sha1 import sha1
 from homeassistant.components.climate import DOMAIN as CLIMATE_DOMAIN
 from homeassistant.components.group import expand_entity_ids
+from homeassistant.components.recorder import get_instance, history
 from homeassistant.components.sensor import STATE_CLASS_MEASUREMENT, SensorEntity
 from homeassistant.components.water_heater import DOMAIN as WATER_HEATER_DOMAIN
 from homeassistant.components.weather import DOMAIN as WEATHER_DOMAIN
@@ -34,7 +37,7 @@
     STATE_UNAVAILABLE,
     STATE_UNKNOWN,
 )
-from homeassistant.core import HomeAssistant, callback, split_entity_id
+from homeassistant.core import HomeAssistant, State, callback, split_entity_id
 from homeassistant.exceptions import TemplateError
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.config_validation import PLATFORM_SCHEMA
@@ -57,16 +60,6 @@
     UPDATE_MIN_TIME,
 )
 
-try:  # pragma: no cover
-    # HA version >=2021.6
-    from homeassistant.components.recorder import get_instance, history
-    from homeassistant.components.recorder.models import LazyState
-except ImportError:  # pragma: no cover
-    # HA version <=2021.5
-    from homeassistant.components import history
-    from homeassistant.components.history import LazyState
-
-
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -248,7 +241,7 @@ def _has_state(state) -> bool:
             "",
         ]
 
-    def _get_temperature(self, state: LazyState) -> Optional[float]:
+    def _get_temperature(self, state: State) -> Optional[float]:
         """Get temperature value from entity."""
         ha_unit = self.hass.config.units.temperature_unit
         domain = split_entity_id(state.entity_id)[0]
@@ -273,7 +266,7 @@ def _get_temperature(self, state: LazyState) -> Optional[float]:
 
         return temperature
 
-    def _get_state_value(self, state: LazyState) -> Optional[float]:
+    def _get_state_value(self, state: State) -> Optional[float]:
         """Return value of given entity state and count some sensor attributes."""
         state = self._get_temperature(state) if self._temperature_mode else state.state
         if not self._has_state(state):
@@ -387,7 +380,7 @@ async def _async_update_period(self):  # pylint: disable=too-many-branches
         self.start = start.replace(microsecond=0).isoformat()
         self.end = end.replace(microsecond=0).isoformat()
 
-    def _init_mode(self, state: LazyState):
+    def _init_mode(self, state: State):
         """Initialize sensor mode."""
         if self._temperature_mode is not None:
             return
@@ -458,7 +451,7 @@ async def _async_update_state(
         for entity_id in self.sources:
             _LOGGER.debug('Processing entity "%s"', entity_id)
 
-            state = self.hass.states.get(entity_id)  # type: LazyState
+            state = self.hass.states.get(entity_id)  # type: State
 
             if state is None:
                 _LOGGER.error('Unable to find an entity "%s"', entity_id)
@@ -539,4 +532,9 @@ async def _async_update_state(
                 self._attr_native_value = int(self._attr_native_value)
         else:
             self._attr_native_value = None
-        _LOGGER.debug("Total average state: %s", self._attr_native_value)
+
+        _LOGGER.debug(
+            "Total average state: %s %s",
+            self._attr_native_value,
+            self._attr_native_unit_of_measurement,
+        )

pylintrc

@@ -16,7 +16,6 @@ good-names=id,i,j,k,ex,Run,_,fp,T
 # locally-disabled - it spams too much
 # duplicate-code - unavoidable
 # cyclic-import - doesn't test if both import on load
-# abstract-class-little-used - prevents from setting right foundation
 # unused-argument - generic callbacks and setup methods create a lot of warnings
 # too-many-* - are not enforced for the sake of readability
 # too-few-* - same as too-many-*
@@ -26,7 +25,6 @@ good-names=id,i,j,k,ex,Run,_,fp,T
 # wrong-import-order - isort guards this
 disable=
   format,
-  abstract-class-little-used,
   abstract-method,
   cyclic-import,
   duplicate-code,

tests/test_recorder.py

@@ -0,0 +1,85 @@
+"""The tests for recorder platform."""
+from __future__ import annotations
+
+import logging
+from datetime import timedelta
+from unittest.mock import patch
+
+from homeassistant.components import recorder
+from homeassistant.components.input_boolean import DOMAIN
+from homeassistant.components.recorder.models import StateAttributes, States
+from homeassistant.components.recorder.util import session_scope
+from homeassistant.const import ATTR_EDITABLE
+from homeassistant.core import HomeAssistant, State
+from homeassistant.setup import async_setup_component
+from homeassistant.util import dt as dt_util
+from pytest_homeassistant_custom_component.common import async_fire_time_changed
+from pytest_homeassistant_custom_component.components.recorder.common import (
+    wait_recording_done,
+)
+
+from custom_components.average.const import ATTR_SOURCES, ATTR_TO_PROPERTY
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_init_recorder_component(hass, add_config=None):
+    """Initialize the recorder asynchronously."""
+    config = dict(add_config) if add_config else {}
+    if recorder.CONF_DB_URL not in config:
+        config[recorder.CONF_DB_URL] = "sqlite://"  # In memory DB
+        if recorder.CONF_COMMIT_INTERVAL not in config:
+            config[recorder.CONF_COMMIT_INTERVAL] = 0
+
+    with patch(
+        "homeassistant.components.recorder.ALLOW_IN_MEMORY_DB",
+        True,
+    ), patch("homeassistant.components.recorder.migration.migrate_schema"):
+        assert await async_setup_component(
+            hass, recorder.DOMAIN, {recorder.DOMAIN: config}
+        )
+        assert recorder.DOMAIN in hass.config.components
+    _LOGGER.info(
+        "Test recorder successfully started, database location: %s",
+        config[recorder.CONF_DB_URL],
+    )
+
+
+async def async_wait_recording_done_without_instance(hass: HomeAssistant) -> None:
+    """Block till recording is done."""
+    await hass.loop.run_in_executor(None, wait_recording_done, hass)
+
+
+async def test_exclude_attributes(
+    hass: HomeAssistant, enable_custom_integrations: None
+):
+    """Test attributes to be excluded."""
+    await async_init_recorder_component(hass)
+    assert await async_setup_component(hass, DOMAIN, {DOMAIN: {"test": {}}})
+
+    state = hass.states.get("input_boolean.test")
+    assert state
+    assert state.attributes[ATTR_EDITABLE] is False
+
+    await hass.async_block_till_done()
+    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(minutes=5))
+    await hass.async_block_till_done()
+    await async_wait_recording_done_without_instance(hass)
+
+    def _fetch_states() -> list[State]:
+        with session_scope(hass=hass) as session:
+            native_states = []
+            for db_state, db_state_attributes in session.query(States, StateAttributes):
+                state = db_state.to_native()
+                state.attributes = db_state_attributes.to_native()
+                native_states.append(state)
+            return native_states
+
+    states: list[State] = await hass.async_add_executor_job(_fetch_states)
+    assert len(states) == 1
+
+    attributes_to_exclude = set(ATTR_TO_PROPERTY)
+    attributes_to_exclude.discard(ATTR_SOURCES)
+
+    for attr in attributes_to_exclude:
+        assert attr not in states[0].attributes