Changed setup of EnergyZero services

homeassistant/components/energyzero/__init__.py

@@ -5,12 +5,23 @@
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.helpers import config_validation as cv
+from homeassistant.helpers.typing import ConfigType
 
 from .const import DOMAIN
 from .coordinator import EnergyZeroDataUpdateCoordinator
-from .services import async_register_services
+from .services import async_setup_services
 
 PLATFORMS = [Platform.SENSOR]
+CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
+
+
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+    """Set up EnergyZero services."""
+
+    await async_setup_services(hass)
+
+    return True
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
@@ -27,8 +38,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
-    async_register_services(hass, coordinator)
-
     return True
 
 

homeassistant/components/energyzero/services.py

@@ -9,19 +9,21 @@
 from energyzero import Electricity, Gas, VatOption
 import voluptuous as vol
 
+from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.core import (
     HomeAssistant,
     ServiceCall,
     ServiceResponse,
     SupportsResponse,
-    callback,
 )
-from homeassistant.exceptions import ServiceValidationError
+from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
+from homeassistant.helpers import selector
 from homeassistant.util import dt as dt_util
 
 from .const import DOMAIN
 from .coordinator import EnergyZeroDataUpdateCoordinator
 
+ATTR_CONFIG_ENTRY: Final = "config_entry"
 ATTR_START: Final = "start"
 ATTR_END: Final = "end"
 ATTR_INCL_VAT: Final = "incl_vat"
@@ -30,6 +32,11 @@
 ENERGY_SERVICE_NAME: Final = "get_energy_prices"
 SERVICE_SCHEMA: Final = vol.Schema(
     {
+        vol.Required(ATTR_CONFIG_ENTRY): selector.ConfigEntrySelector(
+            {
+                "integration": DOMAIN,
+            }
+        ),
         vol.Required(ATTR_INCL_VAT): bool,
         vol.Optional(ATTR_START): str,
         vol.Optional(ATTR_END): str,
@@ -75,12 +82,29 @@
     }
 
 
+async def __get_coordinator(
+    hass: HomeAssistant, call: ServiceCall
+) -> EnergyZeroDataUpdateCoordinator:
+    entry_id: str = call.data[ATTR_CONFIG_ENTRY]
+    entry: ConfigEntry | None = hass.config_entries.async_get_entry(entry_id)
+
+    if not entry:
+        raise HomeAssistantError(f"Invalid config entry: {entry_id}")
+    if entry.state != ConfigEntryState.LOADED:
+        raise HomeAssistantError(f"{entry.title} is not loaded")
+    if not (energyzero_domain_data := hass.data[DOMAIN].get(entry_id)):
+        raise HomeAssistantError(f"Config entry not loaded: {entry_id}")
+    return energyzero_domain_data
+
+
 async def __get_prices(
     call: ServiceCall,
     *,
-    coordinator: EnergyZeroDataUpdateCoordinator,
+    hass: HomeAssistant,
     price_type: PriceType,
 ) -> ServiceResponse:
+    coordinator = await __get_coordinator(hass, call)
+
     start = __get_date(call.data.get(ATTR_START))
     end = __get_date(call.data.get(ATTR_END))
 
@@ -107,23 +131,20 @@
     return __serialize_prices(data)
 
 
-@callback
-def async_register_services(
-    hass: HomeAssistant, coordinator: EnergyZeroDataUpdateCoordinator
-):
+async def async_setup_services(hass: HomeAssistant) -> None:
     """Set up EnergyZero services."""
 
     hass.services.async_register(
         DOMAIN,
         GAS_SERVICE_NAME,
-        partial(__get_prices, coordinator=coordinator, price_type=PriceType.GAS),
+        partial(__get_prices, hass=hass, price_type=PriceType.GAS),
         schema=SERVICE_SCHEMA,
         supports_response=SupportsResponse.ONLY,
     )
     hass.services.async_register(
         DOMAIN,
         ENERGY_SERVICE_NAME,
-        partial(__get_prices, coordinator=coordinator, price_type=PriceType.ENERGY),
+        partial(__get_prices, hass=hass, price_type=PriceType.ENERGY),
         schema=SERVICE_SCHEMA,
         supports_response=SupportsResponse.ONLY,
     )

homeassistant/components/energyzero/services.yaml

@@ -1,5 +1,10 @@
 get_gas_prices:
   fields:
+    config_entry:
+      required: true
+      selector:
+        config_entry:
+          integration: energyzero
     incl_vat:
       required: true
       default: true
@@ -17,6 +22,11 @@ get_gas_prices:
         datetime:
 get_energy_prices:
   fields:
+    config_entry:
+      required: true
+      selector:
+        config_entry:
+          integration: energyzero
     incl_vat:
       required: true
       default: true

homeassistant/components/energyzero/strings.json

@@ -50,6 +50,10 @@
       "name": "Get gas prices",
       "description": "Request gas prices from EnergyZero.",
       "fields": {
+        "config_entry": {
+          "name": "Config Entry",
+          "description": "The config entry to use for this service."
+        },
         "incl_vat": {
           "name": "Including VAT",
           "description": "Include VAT in the prices."
@@ -68,6 +72,10 @@
       "name": "Get energy prices",
       "description": "Request energy prices from EnergyZero.",
       "fields": {
+        "config_entry": {
+          "name": "[%key:component::energyzero::services::get_gas_prices::fields::config_entry::name%]",
+          "description": "[%key:component::energyzero::services::get_gas_prices::fields::config_entry::description%]"
+        },
         "incl_vat": {
           "name": "[%key:component::energyzero::services::get_gas_prices::fields::incl_vat::name%]",
           "description": "[%key:component::energyzero::services::get_gas_prices::fields::incl_vat::description%]"

tests/components/energyzero/test_services.py

@@ -6,12 +6,15 @@
 
 from homeassistant.components.energyzero.const import DOMAIN
 from homeassistant.components.energyzero.services import (
+    ATTR_CONFIG_ENTRY,
     ENERGY_SERVICE_NAME,
     GAS_SERVICE_NAME,
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ServiceValidationError
 
+from tests.common import MockConfigEntry
+
 
 @pytest.mark.usefixtures("init_integration")
 async def test_has_services(
@@ -29,15 +32,17 @@ async def test_has_services(
 @pytest.mark.parametrize("end", [{"end": "2023-01-01 00:00:00"}, {}])
 async def test_service(
     hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
     snapshot: SnapshotAssertion,
     service: str,
     incl_vat: dict[str, bool],
     start: dict[str, str],
     end: dict[str, str],
 ) -> None:
     """Test the EnergyZero Service."""
+    entry = {ATTR_CONFIG_ENTRY: mock_config_entry.entry_id}
 
-    data = incl_vat | start | end
+    data = entry | incl_vat | start | end
 
     assert snapshot == await hass.services.async_call(
         DOMAIN,
@@ -51,20 +56,24 @@ async def test_service(
 @pytest.mark.usefixtures("init_integration")
 @pytest.mark.parametrize("service", [GAS_SERVICE_NAME, ENERGY_SERVICE_NAME])
 @pytest.mark.parametrize(
-    ("service_data", "error", "error_message"),
+    ("include_mock_entry", "service_data", "error", "error_message"),
     [
-        ({}, vol.er.Error, "required key not provided .+"),
+        (False, {}, vol.er.Error, "required key not provided .+"),
+        (True, {}, vol.er.Error, "required key not provided .+"),
         (
+            True,
             {"incl_vat": "incorrect vat"},
             vol.er.Error,
             "expected bool for dictionary value .+",
         ),
         (
+            True,
             {"incl_vat": True, "start": "incorrect date"},
             ServiceValidationError,
             "Invalid datetime provided.",
         ),
         (
+            True,
             {"incl_vat": True, "end": "incorrect date"},
             ServiceValidationError,
             "Invalid datetime provided.",
@@ -73,13 +82,18 @@ async def test_service(
 )
 async def test_service_validation(
     hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    include_mock_entry: bool,
     service: str,
     service_data: dict[str, str],
     error: type[Exception],
     error_message: str,
 ) -> None:
     """Test the EnergyZero Service validation."""
 
+    if include_mock_entry:
+        service_data = service_data | {ATTR_CONFIG_ENTRY: mock_config_entry.entry_id}
+
     with pytest.raises(error, match=error_message):
         await hass.services.async_call(
             DOMAIN,