Isolate aiohttp session and storage per config entry (multi-account support)

custom_components/boilerjuice/__init__.py

@@ -9,12 +9,14 @@
 from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
 from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, Platform
 from homeassistant.core import HomeAssistant, ServiceCall, callback
-from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
+from homeassistant.exceptions import (
+    ConfigEntryAuthFailed,
+    ConfigEntryNotReady,
+    HomeAssistantError,
+)
 from homeassistant.helpers import config_validation as cv
-from homeassistant.helpers import service
-from homeassistant.helpers.device_registry import DeviceEntryType, DeviceRegistry
+from homeassistant.helpers.device_registry import DeviceEntryType
 from homeassistant.helpers.device_registry import async_get as async_get_device_registry
-from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.typing import ConfigType
 
 from .const import CONF_KWH_PER_LITRE, CONF_TANK_ID, DEFAULT_KWH_PER_LITRE, DOMAIN
@@ -40,16 +42,76 @@
 # Service schemas
 SERVICE_RESET_CONSUMPTION = "reset_consumption"
 SERVICE_SET_CONSUMPTION = "set_consumption"
-RESET_CONSUMPTION_SCHEMA = vol.Schema({})
+
+# Optional target selectors so users with multiple BoilerJuice accounts can
+# address a single tank instead of fanning the service call out to every
+# configured entry. HA injects device_id/area_id/entity_id/label_id when the
+# user picks a target from the UI, so allow extras through the voluptuous
+# schema rather than trying to enumerate every key.
+RESET_CONSUMPTION_SCHEMA = vol.Schema(
+    {
+        vol.Optional("entry_id"): vol.Any(cv.string, [cv.string]),
+    },
+    extra=vol.ALLOW_EXTRA,
+)
 
 SET_CONSUMPTION_SCHEMA = vol.Schema(
     {
         vol.Required("liters"): cv.positive_float,
         vol.Optional("daily"): cv.positive_float,
-    }
+        vol.Optional("entry_id"): vol.Any(cv.string, [cv.string]),
+    },
+    extra=vol.ALLOW_EXTRA,
 )
 
 
+def _resolve_target_coordinators(hass: HomeAssistant, call: ServiceCall) -> list:
+    """Return the coordinators a service call should operate on.
+
+    Honours the optional ``device_id`` / ``entry_id`` fields plus any target
+    selector the user picks in the UI. Falls back to all configured entries
+    for backwards compatibility.
+    """
+    entry_ids: set[str] = set()
+
+    def _collect(value):
+        if value is None:
+            return
+        if isinstance(value, str):
+            entry_ids.add(value)
+        else:
+            entry_ids.update(value)
+
+    _collect(call.data.get("entry_id"))
+
+    device_registry = async_get_device_registry(hass)
+    device_ids = call.data.get("device_id")
+    if isinstance(device_ids, str):
+        device_ids = [device_ids]
+    for device_id in device_ids or []:
+        device = device_registry.async_get(device_id)
+        if device is None:
+            raise HomeAssistantError(f"Unknown device_id {device_id}")
+        for entry_id in device.config_entries:
+            if entry_id in hass.data.get(DOMAIN, {}):
+                entry_ids.add(entry_id)
+
+    coordinators_by_entry = hass.data.get(DOMAIN, {})
+
+    if not entry_ids:
+        return list(coordinators_by_entry.values())
+
+    resolved = []
+    for entry_id in entry_ids:
+        coordinator = coordinators_by_entry.get(entry_id)
+        if coordinator is None:
+            raise HomeAssistantError(
+                f"No BoilerJuice integration loaded for entry_id {entry_id}"
+            )
+        resolved.append(coordinator)
+    return resolved
+
+
 @callback
 def async_setup_services(hass: HomeAssistant) -> None:
     """Set up the BoilerJuice services."""
@@ -58,7 +120,7 @@ def async_setup_services(hass: HomeAssistant) -> None:
 
     async def async_handle_reset_consumption(call: ServiceCall) -> None:
         """Handle the service call to reset consumption."""
-        for entry_id, coordinator in hass.data[DOMAIN].items():
+        for coordinator in _resolve_target_coordinators(hass, call):
             coordinator.reset_consumption()
             await coordinator.async_request_refresh()
 
@@ -75,7 +137,7 @@ async def async_handle_set_consumption(call: ServiceCall) -> None:
         total_consumption = data["liters"]
         daily_consumption = data.get("daily")
 
-        for entry_id, coordinator in hass.data[DOMAIN].items():
+        for coordinator in _resolve_target_coordinators(hass, call):
             if coordinator.data:
                 # Set the consumption values
                 coordinator._total_consumption_usable_liters = total_consumption
@@ -125,6 +187,7 @@ def async_unload_services(hass: HomeAssistant) -> None:
         return
 
     hass.services.async_remove(DOMAIN, SERVICE_RESET_CONSUMPTION)
+    hass.services.async_remove(DOMAIN, SERVICE_SET_CONSUMPTION)
 
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
@@ -187,7 +250,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id)
+        coordinator = hass.data[DOMAIN].pop(entry.entry_id, None)
+        if coordinator is not None:
+            await coordinator.async_close()
         if not hass.data[DOMAIN]:
             async_unload_services(hass)
             hass.data.pop(DOMAIN)

custom_components/boilerjuice/config_flow.py

@@ -35,23 +35,28 @@ async def validate_input(hass: HomeAssistant, data: dict[str, Any]) -> dict[str,
     coordinator = BoilerJuiceDataUpdateCoordinator(hass, data)
 
     try:
-        await coordinator.async_refresh()
-    except Exception as err:
-        if "Invalid credentials" in str(err):
-            raise InvalidAuth from err
-        if "Failed to login" in str(err):
-            raise CannotConnect from err
-        raise err
-
-    # Get the model name if available, fallback to tank name, then default
-    title = "BoilerJuice Tank"
-    if coordinator.data:
-        if coordinator.data.get("model"):
-            title = coordinator.data["model"]
-        elif coordinator.data.get("name"):
-            title = coordinator.data["name"]
-
-    return {"title": title}
+        try:
+            await coordinator.async_refresh()
+        except Exception as err:
+            if "Invalid credentials" in str(err):
+                raise InvalidAuth from err
+            if "Failed to login" in str(err):
+                raise CannotConnect from err
+            raise err
+
+        # Get the model name if available, fallback to tank name, then default
+        title = "BoilerJuice Tank"
+        if coordinator.data:
+            if coordinator.data.get("model"):
+                title = coordinator.data["model"]
+            elif coordinator.data.get("name"):
+                title = coordinator.data["name"]
+
+        return {"title": title}
+    finally:
+        # Each coordinator owns a private aiohttp session; always close the
+        # throwaway one used for config-flow validation.
+        await coordinator.async_close()
 
 
 class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):