Merge pull request #122 from danieldotnl/refactor_reset

Refactor reset service implementation
danieldotnl · Apr 19, 2024 · dd5a1d7 · dd5a1d7
2 parents 3298e71 + 9b96059
commit dd5a1d7
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 70 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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,
     ):