Modify RainMachine to store a single dataclass in hass.data (#75460)

* Modify RainMachine to store a single dataclass in `hass.data`

* Pass one object around instead of multiple

homeassistant/components/rainmachine/__init__.py

@@ -2,9 +2,10 @@
 from __future__ import annotations
 
 import asyncio
+from dataclasses import dataclass
 from datetime import timedelta
 from functools import partial
-from typing import Any, cast
+from typing import Any
 
 from regenmaschine import Client
 from regenmaschine.controller import Controller
@@ -28,7 +29,7 @@ from homeassistant.helpers import (
     device_registry as dr,
     entity_registry as er,
 )
-from homeassistant.helpers.entity import DeviceInfo, EntityDescription
+from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.update_coordinator import (
     CoordinatorEntity,
     DataUpdateCoordinator,
@@ -39,8 +40,6 @@ from homeassistant.util.network import is_ip_address
 from .config_flow import get_client_controller
 from .const import (
     CONF_ZONE_RUN_TIME,
-    DATA_CONTROLLER,
-    DATA_COORDINATOR,
     DATA_PROGRAMS,
     DATA_PROVISION_SETTINGS,
     DATA_RESTRICTIONS_CURRENT,
@@ -49,6 +48,7 @@ from .const import (
     DOMAIN,
     LOGGER,
 )
+from .model import RainMachineEntityDescription
 
 DEFAULT_SSL = True
 
@@ -135,6 +135,14 @@ SERVICE_RESTRICT_WATERING_SCHEMA = SERVICE_SCHEMA.extend(
 )
 
 
+@dataclass
+class RainMachineData:
+    """Define an object to be stored in `hass.data`."""
+
+    controller: Controller
+    coordinators: dict[str, DataUpdateCoordinator]
+
+
 @callback
 def async_get_controller_for_service_call(
     hass: HomeAssistant, call: ServiceCall
@@ -146,9 +154,8 @@ def async_get_controller_for_service_call(
     if device_entry := device_registry.async_get(device_id):
         for entry in hass.config_entries.async_entries(DOMAIN):
             if entry.entry_id in device_entry.config_entries:
-                return cast(
+                data: RainMachineData = hass.data[DOMAIN][entry.entry_id]
-                    Controller, hass.data[DOMAIN][entry.entry_id][DATA_CONTROLLER]
+                return data.controller
-                )
 
     raise ValueError(f"No controller for device ID: {device_id}")
 
@@ -161,14 +168,12 @@ async def async_update_programs_and_zones(
     Program and zone updates always go together because of how linked they are:
     programs affect zones and certain combinations of zones affect programs.
     """
+    data: RainMachineData = hass.data[DOMAIN][entry.entry_id]
+
     await asyncio.gather(
         *[
-            hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR][
+            data.coordinators[DATA_PROGRAMS].async_refresh(),
-                DATA_PROGRAMS
+            data.coordinators[DATA_ZONES].async_refresh(),
-            ].async_refresh(),
-            hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR][
-                DATA_ZONES
-            ].async_refresh(),
         ]
     )
 
@@ -250,10 +255,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     await asyncio.gather(*controller_init_tasks)
 
     hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = {
+    hass.data[DOMAIN][entry.entry_id] = RainMachineData(
-        DATA_CONTROLLER: controller,
+        controller=controller, coordinators=coordinators
-        DATA_COORDINATOR: coordinators,
+    )
-    }
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -406,28 +410,27 @@ class RainMachineEntity(CoordinatorEntity):
     def __init__(
         self,
         entry: ConfigEntry,
-        coordinator: DataUpdateCoordinator,
+        data: RainMachineData,
-        controller: Controller,
+        description: RainMachineEntityDescription,
-        description: EntityDescription,
     ) -> None:
         """Initialize."""
-        super().__init__(coordinator)
+        super().__init__(data.coordinators[description.api_category])
 
         self._attr_device_info = DeviceInfo(
-            identifiers={(DOMAIN, controller.mac)},
+            identifiers={(DOMAIN, data.controller.mac)},
             configuration_url=f"https://{entry.data[CONF_IP_ADDRESS]}:{entry.data[CONF_PORT]}",
-            connections={(dr.CONNECTION_NETWORK_MAC, controller.mac)},
+            connections={(dr.CONNECTION_NETWORK_MAC, data.controller.mac)},
-            name=str(controller.name).capitalize(),
+            name=str(data.controller.name).capitalize(),
             manufacturer="RainMachine",
             model=(
-                f"Version {controller.hardware_version} "
+                f"Version {data.controller.hardware_version} "
-                f"(API: {controller.api_version})"
+                f"(API: {data.controller.api_version})"
             ),
-            sw_version=controller.software_version,
+            sw_version=data.controller.software_version,
         )
         self._attr_extra_state_attributes = {}
-        self._attr_unique_id = f"{controller.mac}_{description.key}"
+        self._attr_unique_id = f"{data.controller.mac}_{description.key}"
-        self._controller = controller
+        self._data = data
         self.entity_description = description
 
     @callback

homeassistant/components/rainmachine/binary_sensor.py

@@ -10,16 +10,17 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import RainMachineEntity
+from . import RainMachineData, RainMachineEntity
 from .const import (
-    DATA_CONTROLLER,
-    DATA_COORDINATOR,
     DATA_PROVISION_SETTINGS,
     DATA_RESTRICTIONS_CURRENT,
     DATA_RESTRICTIONS_UNIVERSAL,
     DOMAIN,
 )
-from .model import RainMachineDescriptionMixinApiCategory
+from .model import (
+    RainMachineEntityDescription,
+    RainMachineEntityDescriptionMixinDataKey,
+)
 from .util import key_exists
 
 TYPE_FLOW_SENSOR = "flow_sensor"
@@ -35,7 +36,9 @@ TYPE_WEEKDAY = "weekday"
 
 @dataclass
 class RainMachineBinarySensorDescription(
-    BinarySensorEntityDescription, RainMachineDescriptionMixinApiCategory
+    BinarySensorEntityDescription,
+    RainMachineEntityDescription,
+    RainMachineEntityDescriptionMixinDataKey,
 ):
     """Describe a RainMachine binary sensor."""
 
@@ -124,8 +127,7 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up RainMachine binary sensors based on a config entry."""
-    controller = hass.data[DOMAIN][entry.entry_id][DATA_CONTROLLER]
+    data: RainMachineData = hass.data[DOMAIN][entry.entry_id]
-    coordinators = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR]
 
     api_category_sensor_map = {
         DATA_PROVISION_SETTINGS: ProvisionSettingsBinarySensor,
@@ -135,12 +137,10 @@ async def async_setup_entry(
 
     async_add_entities(
         [
-            api_category_sensor_map[description.api_category](
+            api_category_sensor_map[description.api_category](entry, data, description)
-                entry, coordinator, controller, description
-            )
             for description in BINARY_SENSOR_DESCRIPTIONS
             if (
-                (coordinator := coordinators[description.api_category]) is not None
+                (coordinator := data.coordinators[description.api_category]) is not None
                 and coordinator.data
                 and key_exists(coordinator.data, description.data_key)
             )
@@ -151,6 +151,8 @@ async def async_setup_entry(
 class CurrentRestrictionsBinarySensor(RainMachineEntity, BinarySensorEntity):
     """Define a binary sensor that handles current restrictions data."""
 
+    entity_description: RainMachineBinarySensorDescription
+
     @callback
     def update_from_latest_data(self) -> None:
         """Update the state."""
@@ -171,6 +173,8 @@ class CurrentRestrictionsBinarySensor(RainMachineEntity, BinarySensorEntity):
 class ProvisionSettingsBinarySensor(RainMachineEntity, BinarySensorEntity):
     """Define a binary sensor that handles provisioning data."""
 
+    entity_description: RainMachineBinarySensorDescription
+
     @callback
     def update_from_latest_data(self) -> None:
         """Update the state."""
@@ -181,6 +185,8 @@ class ProvisionSettingsBinarySensor(RainMachineEntity, BinarySensorEntity):
 class UniversalRestrictionsBinarySensor(RainMachineEntity, BinarySensorEntity):
     """Define a binary sensor that handles universal restrictions data."""
 
+    entity_description: RainMachineBinarySensorDescription
+
     @callback
     def update_from_latest_data(self) -> None:
         """Update the state."""

homeassistant/components/rainmachine/const.py

@@ -7,8 +7,6 @@ DOMAIN = "rainmachine"
 
 CONF_ZONE_RUN_TIME = "zone_run_time"
 
-DATA_CONTROLLER = "controller"
-DATA_COORDINATOR = "coordinator"
 DATA_PROGRAMS = "programs"
 DATA_PROVISION_SETTINGS = "provision.settings"
 DATA_RESTRICTIONS_CURRENT = "restrictions.current"

homeassistant/components/rainmachine/diagnostics.py

@@ -3,15 +3,13 @@ from __future__ import annotations
 
 from typing import Any
 
-from regenmaschine.controller import Controller
-
 from homeassistant.components.diagnostics import async_redact_data
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_LATITUDE, CONF_LONGITUDE, CONF_PASSWORD
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .const import DATA_CONTROLLER, DATA_COORDINATOR, DOMAIN
+from . import RainMachineData
+from .const import DOMAIN
 
 TO_REDACT = {
     CONF_LATITUDE,
@@ -24,9 +22,7 @@ async def async_get_config_entry_diagnostics(
     hass: HomeAssistant, entry: ConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    data = hass.data[DOMAIN][entry.entry_id]
+    data: RainMachineData = hass.data[DOMAIN][entry.entry_id]
-    coordinators: dict[str, DataUpdateCoordinator] = data[DATA_COORDINATOR]
-    controller: Controller = data[DATA_CONTROLLER]
 
     return {
         "entry": {
@@ -38,15 +34,15 @@ async def async_get_config_entry_diagnostics(
             "coordinator": async_redact_data(
                 {
                     api_category: controller.data
-                    for api_category, controller in coordinators.items()
+                    for api_category, controller in data.coordinators.items()
                 },
                 TO_REDACT,
             ),
             "controller": {
-                "api_version": controller.api_version,
+                "api_version": data.controller.api_version,
-                "hardware_version": controller.hardware_version,
+                "hardware_version": data.controller.hardware_version,
-                "name": controller.name,
+                "name": data.controller.name,
-                "software_version": controller.software_version,
+                "software_version": data.controller.software_version,
             },
         },
     }

homeassistant/components/rainmachine/model.py

@@ -1,17 +1,32 @@
 """Define RainMachine data models."""
 from dataclasses import dataclass
 
+from homeassistant.helpers.entity import EntityDescription
+
 
 @dataclass
-class RainMachineDescriptionMixinApiCategory:
+class RainMachineEntityDescriptionMixinApiCategory:
-    """Define an entity description mixin for binary and regular sensors."""
+    """Define an entity description mixin to include an API category."""
 
     api_category: str
+
+
+@dataclass
+class RainMachineEntityDescriptionMixinDataKey:
+    """Define an entity description mixin to include a data payload key."""
+
     data_key: str
 
 
 @dataclass
-class RainMachineDescriptionMixinUid:
+class RainMachineEntityDescriptionMixinUid:
-    """Define an entity description mixin for switches."""
+    """Define an entity description mixin to include an activity UID."""
 
     uid: int
+
+
+@dataclass
+class RainMachineEntityDescription(
+    EntityDescription, RainMachineEntityDescriptionMixinApiCategory
+):
+    """Describe a RainMachine entity."""

homeassistant/components/rainmachine/sensor.py

@@ -5,8 +5,6 @@ from dataclasses import dataclass
 from datetime import datetime, timedelta
 from typing import Any, cast
 
-from regenmaschine.controller import Controller
-
 from homeassistant.components.sensor import (
     RestoreSensor,
     SensorDeviceClass,
@@ -17,15 +15,12 @@ from homeassistant.components.sensor import (
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import TEMP_CELSIUS, VOLUME_CUBIC_METERS
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.entity import EntityCategory, EntityDescription
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 from homeassistant.util.dt import utcnow
 
-from . import RainMachineEntity
+from . import RainMachineData, RainMachineEntity
 from .const import (
-    DATA_CONTROLLER,
-    DATA_COORDINATOR,
     DATA_PROGRAMS,
     DATA_PROVISION_SETTINGS,
     DATA_RESTRICTIONS_UNIVERSAL,
@@ -33,8 +28,9 @@ from .const import (
     DOMAIN,
 )
 from .model import (
-    RainMachineDescriptionMixinApiCategory,
+    RainMachineEntityDescription,
-    RainMachineDescriptionMixinUid,
+    RainMachineEntityDescriptionMixinDataKey,
+    RainMachineEntityDescriptionMixinUid,
 )
 from .util import RUN_STATE_MAP, RunStates, key_exists
 
@@ -50,21 +46,25 @@ TYPE_ZONE_RUN_COMPLETION_TIME = "zone_run_completion_time"
 
 
 @dataclass
-class RainMachineSensorDescriptionApiCategory(
+class RainMachineSensorDataDescription(
-    SensorEntityDescription, RainMachineDescriptionMixinApiCategory
+    SensorEntityDescription,
+    RainMachineEntityDescription,
+    RainMachineEntityDescriptionMixinDataKey,
 ):
     """Describe a RainMachine sensor."""
 
 
 @dataclass
-class RainMachineSensorDescriptionUid(
+class RainMachineSensorCompletionTimerDescription(
-    SensorEntityDescription, RainMachineDescriptionMixinUid
+    SensorEntityDescription,
+    RainMachineEntityDescription,
+    RainMachineEntityDescriptionMixinUid,
 ):
     """Describe a RainMachine sensor."""
 
 
 SENSOR_DESCRIPTIONS = (
-    RainMachineSensorDescriptionApiCategory(
+    RainMachineSensorDataDescription(
         key=TYPE_FLOW_SENSOR_CLICK_M3,
         name="Flow sensor clicks per cubic meter",
         icon="mdi:water-pump",
@@ -75,7 +75,7 @@ SENSOR_DESCRIPTIONS = (
         api_category=DATA_PROVISION_SETTINGS,
         data_key="flowSensorClicksPerCubicMeter",
     ),
-    RainMachineSensorDescriptionApiCategory(
+    RainMachineSensorDataDescription(
         key=TYPE_FLOW_SENSOR_CONSUMED_LITERS,
         name="Flow sensor consumed liters",
         icon="mdi:water-pump",
@@ -86,7 +86,7 @@ SENSOR_DESCRIPTIONS = (
         api_category=DATA_PROVISION_SETTINGS,
         data_key="flowSensorWateringClicks",
     ),
-    RainMachineSensorDescriptionApiCategory(
+    RainMachineSensorDataDescription(
         key=TYPE_FLOW_SENSOR_START_INDEX,
         name="Flow sensor start index",
         icon="mdi:water-pump",
@@ -96,7 +96,7 @@ SENSOR_DESCRIPTIONS = (
         api_category=DATA_PROVISION_SETTINGS,
         data_key="flowSensorStartIndex",
     ),
-    RainMachineSensorDescriptionApiCategory(
+    RainMachineSensorDataDescription(
         key=TYPE_FLOW_SENSOR_WATERING_CLICKS,
         name="Flow sensor clicks",
         icon="mdi:water-pump",
@@ -107,7 +107,7 @@ SENSOR_DESCRIPTIONS = (
         api_category=DATA_PROVISION_SETTINGS,
         data_key="flowSensorWateringClicks",
     ),
-    RainMachineSensorDescriptionApiCategory(
+    RainMachineSensorDataDescription(
         key=TYPE_FREEZE_TEMP,
         name="Freeze protect temperature",
         icon="mdi:thermometer",
@@ -125,8 +125,7 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up RainMachine sensors based on a config entry."""
-    controller = hass.data[DOMAIN][entry.entry_id][DATA_CONTROLLER]
+    data: RainMachineData = hass.data[DOMAIN][entry.entry_id]
-    coordinators = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR]
 
     api_category_sensor_map = {
         DATA_PROVISION_SETTINGS: ProvisionSettingsSensor,
@@ -134,32 +133,29 @@ async def async_setup_entry(
     }
 
     sensors = [
-        api_category_sensor_map[description.api_category](
+        api_category_sensor_map[description.api_category](entry, data, description)
-            entry, coordinator, controller, description
-        )
         for description in SENSOR_DESCRIPTIONS
         if (
-            (coordinator := coordinators[description.api_category]) is not None
+            (coordinator := data.coordinators[description.api_category]) is not None
             and coordinator.data
             and key_exists(coordinator.data, description.data_key)
         )
     ]
 
-    program_coordinator = coordinators[DATA_PROGRAMS]
+    program_coordinator = data.coordinators[DATA_PROGRAMS]
-    zone_coordinator = coordinators[DATA_ZONES]
+    zone_coordinator = data.coordinators[DATA_ZONES]
 
     for uid, program in program_coordinator.data.items():
         sensors.append(
             ProgramTimeRemainingSensor(
                 entry,
-                program_coordinator,
+                data,
-                zone_coordinator,
+                RainMachineSensorCompletionTimerDescription(
-                controller,
-                RainMachineSensorDescriptionUid(
                     key=f"{TYPE_PROGRAM_RUN_COMPLETION_TIME}_{uid}",
                     name=f"{program['name']} Run Completion Time",
                     device_class=SensorDeviceClass.TIMESTAMP,
                     entity_category=EntityCategory.DIAGNOSTIC,
+                    api_category=DATA_PROGRAMS,
                     uid=uid,
                 ),
             )
@@ -169,13 +165,13 @@ async def async_setup_entry(
         sensors.append(
             ZoneTimeRemainingSensor(
                 entry,
-                zone_coordinator,
+                data,
-                controller,
+                RainMachineSensorCompletionTimerDescription(
-                RainMachineSensorDescriptionUid(
                     key=f"{TYPE_ZONE_RUN_COMPLETION_TIME}_{uid}",
                     name=f"{zone['name']} Run Completion Time",
                     device_class=SensorDeviceClass.TIMESTAMP,
                     entity_category=EntityCategory.DIAGNOSTIC,
+                    api_category=DATA_ZONES,
                     uid=uid,
                 ),
             )
@@ -187,17 +183,16 @@ async def async_setup_entry(
 class TimeRemainingSensor(RainMachineEntity, RestoreSensor):
     """Define a sensor that shows the amount of time remaining for an activity."""
 
-    entity_description: RainMachineSensorDescriptionUid
+    entity_description: RainMachineSensorCompletionTimerDescription
 
     def __init__(
         self,
         entry: ConfigEntry,
-        coordinator: DataUpdateCoordinator,
+        data: RainMachineData,
-        controller: Controller,
+        description: RainMachineSensorCompletionTimerDescription,
-        description: EntityDescription,
     ) -> None:
         """Initialize."""
-        super().__init__(entry, coordinator, controller, description)
+        super().__init__(entry, data, description)
 
         self._current_run_state: RunStates | None = None
         self._previous_run_state: RunStates | None = None
@@ -256,19 +251,6 @@ class TimeRemainingSensor(RainMachineEntity, RestoreSensor):
 class ProgramTimeRemainingSensor(TimeRemainingSensor):
     """Define a sensor that shows the amount of time remaining for a program."""
 
-    def __init__(
-        self,
-        entry: ConfigEntry,
-        program_coordinator: DataUpdateCoordinator,
-        zone_coordinator: DataUpdateCoordinator,
-        controller: Controller,
-        description: EntityDescription,
-    ) -> None:
-        """Initialize."""
-        super().__init__(entry, program_coordinator, controller, description)
-
-        self._zone_coordinator = zone_coordinator
-
     @property
     def status_key(self) -> str:
         """Return the data key that contains the activity status."""
@@ -277,7 +259,7 @@ class ProgramTimeRemainingSensor(TimeRemainingSensor):
     def calculate_seconds_remaining(self) -> int:
         """Calculate the number of seconds remaining."""
         return sum(
-            self._zone_coordinator.data[zone["id"]]["remaining"]
+            self._data.coordinators[DATA_ZONES].data[zone["id"]]["remaining"]
             for zone in [z for z in self.activity_data["wateringTimes"] if z["active"]]
         )
 
@@ -285,6 +267,8 @@ class ProgramTimeRemainingSensor(TimeRemainingSensor):
 class ProvisionSettingsSensor(RainMachineEntity, SensorEntity):
     """Define a sensor that handles provisioning data."""
 
+    entity_description: RainMachineSensorDataDescription
+
     @callback
     def update_from_latest_data(self) -> None:
         """Update the state."""
@@ -315,6 +299,8 @@ class ProvisionSettingsSensor(RainMachineEntity, SensorEntity):
 class UniversalRestrictionsSensor(RainMachineEntity, SensorEntity):
     """Define a sensor that handles universal restrictions data."""
 
+    entity_description: RainMachineSensorDataDescription
+
     @callback
     def update_from_latest_data(self) -> None:
         """Update the state."""

homeassistant/components/rainmachine/switch.py

@@ -7,7 +7,6 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import Any
 
-from regenmaschine.controller import Controller
 from regenmaschine.errors import RequestError
 import voluptuous as vol
 
@@ -19,19 +18,16 @@ from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import config_validation as cv, entity_platform
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from . import RainMachineEntity, async_update_programs_and_zones
+from . import RainMachineData, RainMachineEntity, async_update_programs_and_zones
 from .const import (
     CONF_ZONE_RUN_TIME,
-    DATA_CONTROLLER,
-    DATA_COORDINATOR,
     DATA_PROGRAMS,
     DATA_ZONES,
     DEFAULT_ZONE_RUN,
     DOMAIN,
 )
-from .model import RainMachineDescriptionMixinUid
+from .model import RainMachineEntityDescription, RainMachineEntityDescriptionMixinUid
 from .util import RUN_STATE_MAP
 
 ATTR_AREA = "area"
@@ -110,7 +106,9 @@ VEGETATION_MAP = {
 
 @dataclass
 class RainMachineSwitchDescription(
-    SwitchEntityDescription, RainMachineDescriptionMixinUid
+    SwitchEntityDescription,
+    RainMachineEntityDescription,
+    RainMachineEntityDescriptionMixinUid,
 ):
     """Describe a RainMachine switch."""
 
@@ -137,30 +135,27 @@ async def async_setup_entry(
     ):
         platform.async_register_entity_service(service_name, schema, method)
 
-    data = hass.data[DOMAIN][entry.entry_id]
+    data: RainMachineData = hass.data[DOMAIN][entry.entry_id]
-    controller = data[DATA_CONTROLLER]
-    program_coordinator = data[DATA_COORDINATOR][DATA_PROGRAMS]
-    zone_coordinator = data[DATA_COORDINATOR][DATA_ZONES]
 
     entities: list[RainMachineActivitySwitch | RainMachineEnabledSwitch] = []
-
+    for kind, api_category, switch_class, switch_enabled_class in (
-    for kind, coordinator, switch_class, switch_enabled_class in (
+        ("program", DATA_PROGRAMS, RainMachineProgram, RainMachineProgramEnabled),
-        ("program", program_coordinator, RainMachineProgram, RainMachineProgramEnabled),
+        ("zone", DATA_ZONES, RainMachineZone, RainMachineZoneEnabled),
-        ("zone", zone_coordinator, RainMachineZone, RainMachineZoneEnabled),
     ):
-        for uid, data in coordinator.data.items():
+        coordinator = data.coordinators[api_category]
-            name = data["name"].capitalize()
+        for uid, activity in coordinator.data.items():
+            name = activity["name"].capitalize()
 
             # Add a switch to start/stop the program or zone:
             entities.append(
                 switch_class(
                     entry,
-                    coordinator,
+                    data,
-                    controller,
                     RainMachineSwitchDescription(
                         key=f"{kind}_{uid}",
                         name=name,
                         icon="mdi:water",
+                        api_category=api_category,
                         uid=uid,
                     ),
                 )
@@ -170,13 +165,13 @@ async def async_setup_entry(
             entities.append(
                 switch_enabled_class(
                     entry,
-                    coordinator,
+                    data,
-                    controller,
                     RainMachineSwitchDescription(
                         key=f"{kind}_{uid}_enabled",
                         name=f"{name} enabled",
                         entity_category=EntityCategory.CONFIG,
                         icon="mdi:cog",
+                        api_category=api_category,
                         uid=uid,
                     ),
                 )
@@ -193,12 +188,11 @@ class RainMachineBaseSwitch(RainMachineEntity, SwitchEntity):
     def __init__(
         self,
         entry: ConfigEntry,
-        coordinator: DataUpdateCoordinator,
+        data: RainMachineData,
-        controller: Controller,
         description: RainMachineSwitchDescription,
     ) -> None:
         """Initialize."""
-        super().__init__(entry, coordinator, controller, description)
+        super().__init__(entry, data, description)
 
         self._attr_is_on = False
         self._entry = entry
@@ -299,13 +293,13 @@ class RainMachineProgram(RainMachineActivitySwitch):
     async def async_turn_off_when_active(self, **kwargs: Any) -> None:
         """Turn the switch off when its associated activity is active."""
         await self._async_run_api_coroutine(
-            self._controller.programs.stop(self.entity_description.uid)
+            self._data.controller.programs.stop(self.entity_description.uid)
         )
 
     async def async_turn_on_when_active(self, **kwargs: Any) -> None:
         """Turn the switch on when its associated activity is active."""
         await self._async_run_api_coroutine(
-            self._controller.programs.start(self.entity_description.uid)
+            self._data.controller.programs.start(self.entity_description.uid)
         )
 
     @callback
@@ -342,10 +336,10 @@ class RainMachineProgramEnabled(RainMachineEnabledSwitch):
         """Disable the program."""
         tasks = [
             self._async_run_api_coroutine(
-                self._controller.programs.stop(self.entity_description.uid)
+                self._data.controller.programs.stop(self.entity_description.uid)
             ),
             self._async_run_api_coroutine(
-                self._controller.programs.disable(self.entity_description.uid)
+                self._data.controller.programs.disable(self.entity_description.uid)
             ),
         ]
 
@@ -354,7 +348,7 @@ class RainMachineProgramEnabled(RainMachineEnabledSwitch):
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Enable the program."""
         await self._async_run_api_coroutine(
-            self._controller.programs.enable(self.entity_description.uid)
+            self._data.controller.programs.enable(self.entity_description.uid)
         )
 
 
@@ -372,13 +366,13 @@ class RainMachineZone(RainMachineActivitySwitch):
     async def async_turn_off_when_active(self, **kwargs: Any) -> None:
         """Turn the switch off when its associated activity is active."""
         await self._async_run_api_coroutine(
-            self._controller.zones.stop(self.entity_description.uid)
+            self._data.controller.zones.stop(self.entity_description.uid)
         )
 
     async def async_turn_on_when_active(self, **kwargs: Any) -> None:
         """Turn the switch on when its associated activity is active."""
         await self._async_run_api_coroutine(
-            self._controller.zones.start(
+            self._data.controller.zones.start(
                 self.entity_description.uid,
                 kwargs.get("duration", self._entry.options[CONF_ZONE_RUN_TIME]),
             )
@@ -426,10 +420,10 @@ class RainMachineZoneEnabled(RainMachineEnabledSwitch):
         """Disable the zone."""
         tasks = [
             self._async_run_api_coroutine(
-                self._controller.zones.stop(self.entity_description.uid)
+                self._data.controller.zones.stop(self.entity_description.uid)
             ),
             self._async_run_api_coroutine(
-                self._controller.zones.disable(self.entity_description.uid)
+                self._data.controller.zones.disable(self.entity_description.uid)
             ),
         ]
 
@@ -438,5 +432,5 @@ class RainMachineZoneEnabled(RainMachineEnabledSwitch):
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Enable the zone."""
         await self._async_run_api_coroutine(
-            self._controller.zones.enable(self.entity_description.uid)
+            self._data.controller.zones.enable(self.entity_description.uid)
         )