Start using runtime_data for zwave_js (home-assistant#117261)

* Start using runtime_data for zwave_js

* fix bug

homeassistant/components/zwave_js/__init__.py

@@ -182,13 +182,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     # Set up websocket API
     async_register_api(hass)
+    entry.runtime_data = {}
 
     # Create a task to allow the config entry to be unloaded before the driver is ready.
     # Unloading the config entry is needed if the client listen task errors.
     start_client_task = hass.async_create_task(start_client(hass, entry, client))
-    hass.data[DOMAIN].setdefault(entry.entry_id, {})[DATA_START_CLIENT_TASK] = (
-        start_client_task
-    )
+    entry.runtime_data[DATA_START_CLIENT_TASK] = start_client_task
 
     return True
 
@@ -197,9 +196,8 @@ async def start_client(
     hass: HomeAssistant, entry: ConfigEntry, client: ZwaveClient
 ) -> None:
     """Start listening with the client."""
-    entry_hass_data: dict = hass.data[DOMAIN].setdefault(entry.entry_id, {})
-    entry_hass_data[DATA_CLIENT] = client
-    driver_events = entry_hass_data[DATA_DRIVER_EVENTS] = DriverEvents(hass, entry)
+    entry.runtime_data[DATA_CLIENT] = client
+    driver_events = entry.runtime_data[DATA_DRIVER_EVENTS] = DriverEvents(hass, entry)
 
     async def handle_ha_shutdown(event: Event) -> None:
         """Handle HA shutdown."""
@@ -208,7 +206,7 @@ async def handle_ha_shutdown(event: Event) -> None:
     listen_task = asyncio.create_task(
         client_listen(hass, entry, client, driver_events.ready)
     )
-    entry_hass_data[DATA_CLIENT_LISTEN_TASK] = listen_task
+    entry.runtime_data[DATA_CLIENT_LISTEN_TASK] = listen_task
     entry.async_on_unload(
         hass.bus.async_listen(EVENT_HOMEASSISTANT_STOP, handle_ha_shutdown)
     )
@@ -935,11 +933,10 @@ async def client_listen(
 
 async def disconnect_client(hass: HomeAssistant, entry: ConfigEntry) -> None:
     """Disconnect client."""
-    data = hass.data[DOMAIN][entry.entry_id]
-    client: ZwaveClient = data[DATA_CLIENT]
-    listen_task: asyncio.Task = data[DATA_CLIENT_LISTEN_TASK]
-    start_client_task: asyncio.Task = data[DATA_START_CLIENT_TASK]
-    driver_events: DriverEvents = data[DATA_DRIVER_EVENTS]
+    client: ZwaveClient = entry.runtime_data[DATA_CLIENT]
+    listen_task: asyncio.Task = entry.runtime_data[DATA_CLIENT_LISTEN_TASK]
+    start_client_task: asyncio.Task = entry.runtime_data[DATA_START_CLIENT_TASK]
+    driver_events: DriverEvents = entry.runtime_data[DATA_DRIVER_EVENTS]
     listen_task.cancel()
     start_client_task.cancel()
     platform_setup_tasks = driver_events.platform_setup_tasks.values()
@@ -959,9 +956,8 @@ async def disconnect_client(hass: HomeAssistant, entry: ConfigEntry) -> None:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    info = hass.data[DOMAIN][entry.entry_id]
-    client: ZwaveClient = info[DATA_CLIENT]
-    driver_events: DriverEvents = info[DATA_DRIVER_EVENTS]
+    client: ZwaveClient = entry.runtime_data[DATA_CLIENT]
+    driver_events: DriverEvents = entry.runtime_data[DATA_DRIVER_EVENTS]
 
     tasks: list[Coroutine] = [
         hass.config_entries.async_forward_entry_unload(entry, platform)
@@ -973,11 +969,9 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     if client.connected and client.driver:
         await async_disable_server_logging_if_needed(hass, entry, client.driver)
-    if DATA_CLIENT_LISTEN_TASK in info:
+    if DATA_CLIENT_LISTEN_TASK in entry.runtime_data:
         await disconnect_client(hass, entry)
 
-    hass.data[DOMAIN].pop(entry.entry_id)
-
     if entry.data.get(CONF_USE_ADDON) and entry.disabled_by:
         addon_manager: AddonManager = get_addon_manager(hass)
         LOGGER.debug("Stopping Z-Wave JS add-on")
@@ -1016,8 +1010,7 @@ async def async_remove_config_entry_device(
     hass: HomeAssistant, config_entry: ConfigEntry, device_entry: dr.DeviceEntry
 ) -> bool:
     """Remove a config entry from a device."""
-    entry_hass_data = hass.data[DOMAIN][config_entry.entry_id]
-    client: ZwaveClient = entry_hass_data[DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     # Driver may not be ready yet so we can't allow users to remove a device since
     # we need to check if the device is still known to the controller
@@ -1037,7 +1030,7 @@ async def async_remove_config_entry_device(
     ):
         return False
 
-    controller_events: ControllerEvents = entry_hass_data[
+    controller_events: ControllerEvents = config_entry.runtime_data[
         DATA_DRIVER_EVENTS
     ].controller_events
     controller_events.registered_unique_ids.pop(device_entry.id, None)

homeassistant/components/zwave_js/api.py

@@ -75,7 +75,6 @@
 from .const import (
     CONF_DATA_COLLECTION_OPTED_IN,
     DATA_CLIENT,
-    DOMAIN,
     EVENT_DEVICE_ADDED_TO_REGISTRY,
     USER_AGENT,
 )
@@ -285,7 +284,7 @@ async def _async_get_entry(
         )
         return None, None, None
 
-    client: Client = hass.data[DOMAIN][entry_id][DATA_CLIENT]
+    client: Client = entry.runtime_data[DATA_CLIENT]
 
     if client.driver is None:
         connection.send_error(

homeassistant/components/zwave_js/binary_sensor.py

@@ -254,7 +254,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave binary sensor from config entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     @callback
     def async_add_binary_sensor(info: ZwaveDiscoveryInfo) -> None:

homeassistant/components/zwave_js/button.py

@@ -27,7 +27,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave button from config entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     @callback
     def async_add_button(info: ZwaveDiscoveryInfo) -> None:

homeassistant/components/zwave_js/climate.py

@@ -102,7 +102,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave climate from config entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     @callback
     def async_add_climate(info: ZwaveDiscoveryInfo) -> None:

homeassistant/components/zwave_js/cover.py

@@ -57,7 +57,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave Cover from Config Entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     @callback
     def async_add_cover(info: ZwaveDiscoveryInfo) -> None:

homeassistant/components/zwave_js/device_automation_helpers.py

@@ -55,5 +55,5 @@ def async_bypass_dynamic_config_validation(hass: HomeAssistant, device_id: str)
         return True
 
     # The driver may not be ready when the config entry is loaded.
-    client: ZwaveClient = hass.data[DOMAIN][entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = entry.runtime_data[DATA_CLIENT]
     return client.driver is None

homeassistant/components/zwave_js/diagnostics.py

@@ -20,7 +20,7 @@
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import DATA_CLIENT, DOMAIN, USER_AGENT
+from .const import DATA_CLIENT, USER_AGENT
 from .helpers import (
     ZwaveValueMatcher,
     get_home_and_node_id_from_device_entry,
@@ -148,7 +148,7 @@ async def async_get_device_diagnostics(
     hass: HomeAssistant, config_entry: ConfigEntry, device: dr.DeviceEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a device."""
-    client: Client = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: Client = config_entry.runtime_data[DATA_CLIENT]
     identifiers = get_home_and_node_id_from_device_entry(device)
     node_id = identifiers[1] if identifiers else None
     driver = client.driver

homeassistant/components/zwave_js/event.py

@@ -25,7 +25,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave Event entity from Config Entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     @callback
     def async_add_event(info: ZwaveDiscoveryInfo) -> None:

homeassistant/components/zwave_js/fan.py

@@ -49,7 +49,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave Fan from Config Entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     @callback
     def async_add_fan(info: ZwaveDiscoveryInfo) -> None:

homeassistant/components/zwave_js/helpers.py

@@ -155,7 +155,7 @@ async def async_enable_server_logging_if_needed(
     if (curr_server_log_level := driver.log_config.level) and (
         LOG_LEVEL_MAP[curr_server_log_level]
     ) > (lib_log_level := LIB_LOGGER.getEffectiveLevel()):
-        entry_data = hass.data[DOMAIN][entry.entry_id]
+        entry_data = entry.runtime_data
         LOGGER.warning(
             (
                 "Server logging is set to %s and is currently less verbose "
@@ -174,7 +174,6 @@ async def async_disable_server_logging_if_needed(
     hass: HomeAssistant, entry: ConfigEntry, driver: Driver
 ) -> None:
     """Disable logging of zwave-js-server in the lib if still connected to server."""
-    entry_data = hass.data[DOMAIN][entry.entry_id]
     if (
         not driver
         or not driver.client.connected
@@ -183,8 +182,8 @@ async def async_disable_server_logging_if_needed(
         return
     LOGGER.info("Disabling zwave_js server logging")
     if (
-        DATA_OLD_SERVER_LOG_LEVEL in entry_data
-        and (old_server_log_level := entry_data.pop(DATA_OLD_SERVER_LOG_LEVEL))
+        DATA_OLD_SERVER_LOG_LEVEL in entry.runtime_data
+        and (old_server_log_level := entry.runtime_data.pop(DATA_OLD_SERVER_LOG_LEVEL))
         != driver.log_config.level
     ):
         LOGGER.info(
@@ -275,12 +274,12 @@ def async_get_node_from_device_id(
     )
     if entry and entry.state != ConfigEntryState.LOADED:
         raise ValueError(f"Device {device_id} config entry is not loaded")
-    if entry is None or entry.entry_id not in hass.data[DOMAIN]:
+    if entry is None:
         raise ValueError(
             f"Device {device_id} is not from an existing zwave_js config entry"
         )
 
-    client: ZwaveClient = hass.data[DOMAIN][entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = entry.runtime_data[DATA_CLIENT]
     driver = client.driver
 
     if driver is None:
@@ -443,7 +442,9 @@ def async_get_node_status_sensor_entity_id(
     if not (entry_id := _zwave_js_config_entry(hass, device)):
         return None
 
-    client = hass.data[DOMAIN][entry_id][DATA_CLIENT]
+    entry = hass.config_entries.async_get_entry(entry_id)
+    assert entry
+    client = entry.runtime_data[DATA_CLIENT]
     node = async_get_node_from_device_id(hass, device_id, dev_reg)
     return ent_reg.async_get_entity_id(
         SENSOR_DOMAIN,

homeassistant/components/zwave_js/humidifier.py

@@ -73,7 +73,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave humidifier from config entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     @callback
     def async_add_humidifier(info: ZwaveDiscoveryInfo) -> None:

homeassistant/components/zwave_js/light.py

@@ -68,7 +68,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave Light from Config Entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     @callback
     def async_add_light(info: ZwaveDiscoveryInfo) -> None:

homeassistant/components/zwave_js/lock.py

@@ -66,7 +66,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave lock from config entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     @callback
     def async_add_lock(info: ZwaveDiscoveryInfo) -> None:

homeassistant/components/zwave_js/number.py

@@ -31,7 +31,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave Number entity from Config Entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     @callback
     def async_add_number(info: ZwaveDiscoveryInfo) -> None:

homeassistant/components/zwave_js/select.py

@@ -30,7 +30,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave Select entity from Config Entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     @callback
     def async_add_select(info: ZwaveDiscoveryInfo) -> None:

homeassistant/components/zwave_js/sensor.py

@@ -530,7 +530,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave sensor from config entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
     driver = client.driver
     assert driver is not None  # Driver is ready before platforms are loaded.
 

homeassistant/components/zwave_js/services.py

@@ -727,8 +727,8 @@ async def async_multicast_set_value(self, service: ServiceCall) -> None:
             first_node = next(node for node in nodes)
             client = first_node.client
         except StopIteration:
-            entry_id = self._hass.config_entries.async_entries(const.DOMAIN)[0].entry_id
-            client = self._hass.data[const.DOMAIN][entry_id][const.DATA_CLIENT]
+            data = self._hass.config_entries.async_entries(const.DOMAIN)[0].runtime_data
+            client = data[const.DATA_CLIENT]
             assert client.driver
             first_node = next(
                 node

homeassistant/components/zwave_js/siren.py

@@ -33,7 +33,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave Siren entity from Config Entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     @callback
     def async_add_siren(info: ZwaveDiscoveryInfo) -> None:

homeassistant/components/zwave_js/switch.py

@@ -31,7 +31,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave sensor from config entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
 
     @callback
     def async_add_switch(info: ZwaveDiscoveryInfo) -> None:

homeassistant/components/zwave_js/triggers/event.py

@@ -219,7 +219,9 @@ def _create_zwave_listeners() -> None:
         drivers: set[Driver] = set()
         if not (nodes := async_get_nodes_from_targets(hass, config, dev_reg=dev_reg)):
             entry_id = config[ATTR_CONFIG_ENTRY_ID]
-            client: Client = hass.data[DOMAIN][entry_id][DATA_CLIENT]
+            entry = hass.config_entries.async_get_entry(entry_id)
+            assert entry
+            client: Client = entry.runtime_data[DATA_CLIENT]
             driver = client.driver
             assert driver
             drivers.add(driver)

homeassistant/components/zwave_js/triggers/trigger_helpers.py

@@ -37,7 +37,7 @@ def async_bypass_dynamic_config_validation(
             return True
 
         # The driver may not be ready when the config entry is loaded.
-        client: ZwaveClient = hass.data[DOMAIN][entry.entry_id][DATA_CLIENT]
+        client: ZwaveClient = entry.runtime_data[DATA_CLIENT]
         if client.driver is None:
             return True
 

homeassistant/components/zwave_js/update.py

@@ -80,7 +80,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Z-Wave update entity from config entry."""
-    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]
+    client: ZwaveClient = config_entry.runtime_data[DATA_CLIENT]
     cnt: Counter = Counter()
 
     @callback