From 9b96059275bfb9dce115d21f545ecd9e77ce9b59 Mon Sep 17 00:00:00 2001 From: danieldotnl Date: Thu, 18 Apr 2024 14:16:18 +0000 Subject: [PATCH] Refactor implementation of the reset service --- custom_components/measureit/__init__.py | 70 +++---------------------- custom_components/measureit/sensor.py | 21 ++++++++ tests/unit/test_config_flow.py | 3 -- tests/unit/test_options_flow.py | 3 -- 4 files changed, 27 insertions(+), 70 deletions(-) diff --git a/custom_components/measureit/__init__.py b/custom_components/measureit/__init__.py index 92b8b77..8b169d9 100644 --- a/custom_components/measureit/__init__.py +++ b/custom_components/measureit/__init__.py @@ -4,46 +4,18 @@ import voluptuous as vol from homeassistant.config_entries import ConfigEntry -from homeassistant.const import Platform -from homeassistant.const import EVENT_HOMEASSISTANT_STARTED, ATTR_ENTITY_ID -from homeassistant.core import Config, CoreState, callback -from homeassistant.core import HomeAssistant +from homeassistant.const import EVENT_HOMEASSISTANT_STARTED, Platform +from homeassistant.core import CoreState, HomeAssistant, callback from homeassistant.helpers import entity_registry as er from homeassistant.helpers.template import Template -from homeassistant.helpers import config_validation as cv -from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN -from homeassistant.util import dt as dt_util - -from .const import ( - CONF_CONDITION, - CONF_CONFIG_NAME, - CONF_COUNTER_TEMPLATE, - DOMAIN, - EVENT_TYPE_RESET, - MeterType, -) -from .const import CONF_METER_TYPE -from .const import CONF_SOURCE -from .const import CONF_TW_DAYS -from .const import CONF_TW_FROM -from .const import CONF_TW_TILL -from .const import COORDINATOR -from .const import DOMAIN_DATA + +from .const import (CONF_CONDITION, CONF_CONFIG_NAME, CONF_COUNTER_TEMPLATE, + CONF_METER_TYPE, CONF_SOURCE, CONF_TW_DAYS, CONF_TW_FROM, + CONF_TW_TILL, COORDINATOR, DOMAIN_DATA, MeterType) from .coordinator import MeasureItCoordinator from .time_window import TimeWindow _LOGGER: logging.Logger = logging.getLogger(__name__) -CONFIG_SCHEMA = cv.config_entry_only_config_schema( - DOMAIN -) # required to pass hassfest validation due to use of async_setup - - -async def async_setup(hass: HomeAssistant, config: Config): - """Set up this integration using YAML is not supported.""" - hass.data.setdefault(DOMAIN, {}).setdefault(SENSOR_DOMAIN, {}) - _register_services(hass) - return True - async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry): """Set up this integration using UI.""" @@ -117,36 +89,6 @@ def run_start(event): return True -def _register_services(hass: HomeAssistant): - """Register services for MeasureIt.""" - - @callback - def reset_sensor(service_call): - """Reset sensor.""" - _LOGGER.debug("Reset sensor with: %s", service_call.data) - reset_datetime = service_call.data.get("reset_datetime") or dt_util.now() - if not reset_datetime.tzinfo: - reset_datetime = reset_datetime.replace(tzinfo=dt_util.DEFAULT_TIME_ZONE) - - entity_ids = service_call.data[ATTR_ENTITY_ID] - hass.bus.async_fire( - EVENT_TYPE_RESET, - {ATTR_ENTITY_ID: entity_ids, "reset_datetime": reset_datetime}, - ) - - hass.services.async_register( - DOMAIN, - "reset", - reset_sensor, - vol.Schema( - { - vol.Required(ATTR_ENTITY_ID): cv.entity_ids, - vol.Optional("reset_datetime"): cv.datetime, - } - ), - ) - - async def async_reload_entry(hass: HomeAssistant, entry: ConfigEntry) -> None: """Update listener, called when the config entry options are changed.""" await hass.config_entries.async_reload(entry.entry_id) diff --git a/custom_components/measureit/sensor.py b/custom_components/measureit/sensor.py index 96cb2f5..f7b6226 100644 --- a/custom_components/measureit/sensor.py +++ b/custom_components/measureit/sensor.py @@ -102,6 +102,17 @@ async def async_setup_entry( "calibrate", ) + platform.async_register_entity_service( + "reset", + vol.Schema( + { + vol.Required(ATTR_ENTITY_ID): cv.entity_ids, + vol.Optional("reset_datetime"): cv.datetime, + } + ), + "on_reset_service_triggered" + ) + def temp_parse_timestamp_or_string(timestamp_or_string: str) -> datetime | None: """Parse a timestamp or string into a datetime object.""" @@ -376,6 +387,16 @@ def reset(self, event=None): self.schedule_next_reset() self._async_write_ha_state() + @callback + def on_reset_service_triggered(self, reset_datetime: datetime|None = None): + """Handle a reset service call.""" + _LOGGER.debug("Reset sensor with: %s", reset_datetime) + if reset_datetime is None: + reset_datetime = dt_util.now() + if not reset_datetime.tzinfo: + reset_datetime = reset_datetime.replace(tzinfo=dt_util.DEFAULT_TIME_ZONE) + self.schedule_next_reset(reset_datetime) + @callback def schedule_next_reset(self, next_reset: datetime | None = None): """Set the next reset moment.""" diff --git a/tests/unit/test_config_flow.py b/tests/unit/test_config_flow.py index 23882c8..894d37f 100644 --- a/tests/unit/test_config_flow.py +++ b/tests/unit/test_config_flow.py @@ -23,9 +23,6 @@ def bypass_setup_fixture(): """Prevent setup.""" with patch( - "custom_components.measureit.async_setup", - return_value=True, - ), patch( "custom_components.measureit.async_setup_entry", return_value=True, ): diff --git a/tests/unit/test_options_flow.py b/tests/unit/test_options_flow.py index 375480a..4ba0c5d 100644 --- a/tests/unit/test_options_flow.py +++ b/tests/unit/test_options_flow.py @@ -20,9 +20,6 @@ def bypass_setup_fixture(): """Prevent setup.""" with patch( - "custom_components.measureit.async_setup", - return_value=True, - ), patch( "custom_components.measureit.async_setup_entry", return_value=True, ):