Store runtime data inside the config entry in NUT (#116771)

* store runtime data inside the config entry

* remove unsued constants

* add test for InvalidDeviceAutomationConfig exception

* assert entry

* add more specific type hint

homeassistant/components/nut/__init__.py

@@ -26,22 +26,30 @@ from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import (
-    COORDINATOR,
     DEFAULT_SCAN_INTERVAL,
     DOMAIN,
     INTEGRATION_SUPPORTED_COMMANDS,
     PLATFORMS,
-    PYNUT_DATA,
-    PYNUT_UNIQUE_ID,
-    USER_AVAILABLE_COMMANDS,
 )
 
 NUT_FAKE_SERIAL = ["unknown", "blank"]
 
 _LOGGER = logging.getLogger(__name__)
 
+NutConfigEntry = ConfigEntry["NutRuntimeData"]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+@dataclass
+class NutRuntimeData:
+    """Runtime data definition."""
+
+    coordinator: DataUpdateCoordinator
+    data: PyNUTData
+    unique_id: str
+    user_available_commands: set[str]
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: NutConfigEntry) -> bool:
     """Set up Network UPS Tools (NUT) from a config entry."""
 
     # strip out the stale options CONF_RESOURCES,
@@ -110,13 +118,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     else:
         user_available_commands = set()
 
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = {
-        COORDINATOR: coordinator,
-        PYNUT_DATA: data,
-        PYNUT_UNIQUE_ID: unique_id,
-        USER_AVAILABLE_COMMANDS: user_available_commands,
-    }
+    entry.runtime_data = NutRuntimeData(
+        coordinator, data, unique_id, user_available_commands
+    )
 
     device_registry = dr.async_get(hass)
     device_registry.async_get_or_create(
@@ -135,9 +139,7 @@ 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)
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
 async def _async_update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:

homeassistant/components/nut/const.py

@@ -15,15 +15,8 @@ DEFAULT_PORT = 3493
 KEY_STATUS = "ups.status"
 KEY_STATUS_DISPLAY = "ups.status.display"
 
-COORDINATOR = "coordinator"
 DEFAULT_SCAN_INTERVAL = 60
 
-PYNUT_DATA = "data"
-PYNUT_UNIQUE_ID = "unique_id"
-
-
-USER_AVAILABLE_COMMANDS = "user_available_commands"
-
 STATE_TYPES = {
     "OL": "Online",
     "OB": "On Battery",

homeassistant/components/nut/device_action.py

@@ -4,19 +4,15 @@ from __future__ import annotations
 
 import voluptuous as vol
 
+from homeassistant.components.device_automation import InvalidDeviceAutomationConfig
 from homeassistant.const import CONF_DEVICE_ID, CONF_DOMAIN, CONF_TYPE
 from homeassistant.core import Context, HomeAssistant
 from homeassistant.helpers import device_registry as dr
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.typing import ConfigType, TemplateVarsType
 
-from . import PyNUTData
-from .const import (
-    DOMAIN,
-    INTEGRATION_SUPPORTED_COMMANDS,
-    PYNUT_DATA,
-    USER_AVAILABLE_COMMANDS,
-)
+from . import NutRuntimeData
+from .const import DOMAIN, INTEGRATION_SUPPORTED_COMMANDS
 
 ACTION_TYPES = {cmd.replace(".", "_") for cmd in INTEGRATION_SUPPORTED_COMMANDS}
 
@@ -31,18 +27,15 @@ async def async_get_actions(
     hass: HomeAssistant, device_id: str
 ) -> list[dict[str, str]]:
     """List device actions for Network UPS Tools (NUT) devices."""
-    if (entry_id := _get_entry_id_from_device_id(hass, device_id)) is None:
+    if (runtime_data := _get_runtime_data_from_device_id(hass, device_id)) is None:
         return []
     base_action = {
         CONF_DEVICE_ID: device_id,
         CONF_DOMAIN: DOMAIN,
     }
-    user_available_commands: set[str] = hass.data[DOMAIN][entry_id][
-        USER_AVAILABLE_COMMANDS
-    ]
     return [
         {CONF_TYPE: _get_device_action_name(command_name)} | base_action
-        for command_name in user_available_commands
+        for command_name in runtime_data.user_available_commands
     ]
 
 
@@ -56,9 +49,12 @@ async def async_call_action_from_config(
     device_action_name: str = config[CONF_TYPE]
     command_name = _get_command_name(device_action_name)
     device_id: str = config[CONF_DEVICE_ID]
-    entry_id = _get_entry_id_from_device_id(hass, device_id)
-    data: PyNUTData = hass.data[DOMAIN][entry_id][PYNUT_DATA]
-    await data.async_run_command(command_name)
+    runtime_data = _get_runtime_data_from_device_id(hass, device_id)
+    if not runtime_data:
+        raise InvalidDeviceAutomationConfig(
+            f"Unable to find a NUT device with id {device_id}"
+        )
+    await runtime_data.data.async_run_command(command_name)
 
 
 def _get_device_action_name(command_name: str) -> str:
@@ -69,8 +65,14 @@ def _get_command_name(device_action_name: str) -> str:
     return device_action_name.replace("_", ".")
 
 
-def _get_entry_id_from_device_id(hass: HomeAssistant, device_id: str) -> str | None:
+def _get_runtime_data_from_device_id(
+    hass: HomeAssistant, device_id: str
+) -> NutRuntimeData | None:
     device_registry = dr.async_get(hass)
     if (device := device_registry.async_get(device_id)) is None:
         return None
-    return next(entry for entry in device.config_entries)
+    entry = hass.config_entries.async_get_entry(
+        next(entry_id for entry_id in device.config_entries)
+    )
+    assert entry and isinstance(entry.runtime_data, NutRuntimeData)
+    return entry.runtime_data

homeassistant/components/nut/diagnostics.py

@@ -7,27 +7,26 @@ from typing import Any
 import attr
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 
-from . import PyNUTData
-from .const import DOMAIN, PYNUT_DATA, PYNUT_UNIQUE_ID, USER_AVAILABLE_COMMANDS
+from . import NutConfigEntry
+from .const import DOMAIN
 
 TO_REDACT = {CONF_PASSWORD, CONF_USERNAME}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: NutConfigEntry
 ) -> dict[str, dict[str, Any]]:
     """Return diagnostics for a config entry."""
     data = {"entry": async_redact_data(entry.as_dict(), TO_REDACT)}
-    hass_data = hass.data[DOMAIN][entry.entry_id]
+    hass_data = entry.runtime_data
 
     # Get information from Nut library
-    nut_data: PyNUTData = hass_data[PYNUT_DATA]
-    nut_cmd: set[str] = hass_data[USER_AVAILABLE_COMMANDS]
+    nut_data = hass_data.data
+    nut_cmd = hass_data.user_available_commands
     data["nut_data"] = {
         "ups_list": nut_data.ups_list,
         "status": nut_data.status,
@@ -38,7 +37,7 @@ async def async_get_config_entry_diagnostics(
     device_registry = dr.async_get(hass)
     entity_registry = er.async_get(hass)
     hass_device = device_registry.async_get_device(
-        identifiers={(DOMAIN, hass_data[PYNUT_UNIQUE_ID])}
+        identifiers={(DOMAIN, hass_data.unique_id)}
     )
     if not hass_device:
         return data

homeassistant/components/nut/sensor.py

@@ -12,7 +12,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_MANUFACTURER,
     ATTR_MODEL,
@@ -36,16 +35,8 @@ from homeassistant.helpers.update_coordinator import (
     DataUpdateCoordinator,
 )
 
-from . import PyNUTData
-from .const import (
-    COORDINATOR,
-    DOMAIN,
-    KEY_STATUS,
-    KEY_STATUS_DISPLAY,
-    PYNUT_DATA,
-    PYNUT_UNIQUE_ID,
-    STATE_TYPES,
-)
+from . import NutConfigEntry, PyNUTData
+from .const import DOMAIN, KEY_STATUS, KEY_STATUS_DISPLAY, STATE_TYPES
 
 NUT_DEV_INFO_TO_DEV_INFO: dict[str, str] = {
     "manufacturer": ATTR_MANUFACTURER,
@@ -968,15 +959,15 @@ def _get_nut_device_info(data: PyNUTData) -> DeviceInfo:
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: NutConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the NUT sensors."""
 
-    pynut_data = hass.data[DOMAIN][config_entry.entry_id]
-    coordinator = pynut_data[COORDINATOR]
-    data = pynut_data[PYNUT_DATA]
-    unique_id = pynut_data[PYNUT_UNIQUE_ID]
+    pynut_data = config_entry.runtime_data
+    coordinator = pynut_data.coordinator
+    data = pynut_data.data
+    unique_id = pynut_data.unique_id
     status = coordinator.data
 
     resources = [sensor_id for sensor_id in SENSOR_TYPES if sensor_id in status]

tests/components/nut/test_device_action.py

@@ -7,9 +7,13 @@ import pytest
 from pytest_unordered import unordered
 
 from homeassistant.components import automation, device_automation
-from homeassistant.components.device_automation import DeviceAutomationType
+from homeassistant.components.device_automation import (
+    DeviceAutomationType,
+    InvalidDeviceAutomationConfig,
+)
 from homeassistant.components.nut import DOMAIN
 from homeassistant.components.nut.const import INTEGRATION_SUPPORTED_COMMANDS
+from homeassistant.const import CONF_DEVICE_ID, CONF_TYPE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 from homeassistant.setup import async_setup_component
@@ -229,3 +233,25 @@ async def test_rund_command_exception(
     await hass.async_block_till_done()
 
     assert error_message in caplog.text
+
+
+async def test_action_exception_invalid_device(hass: HomeAssistant) -> None:
+    """Test raises exception if invalid device."""
+    list_commands_return_value = {"beeper.enable": None}
+    await async_init_integration(
+        hass,
+        list_vars={"ups.status": "OL"},
+        list_commands_return_value=list_commands_return_value,
+    )
+
+    platform = await device_automation.async_get_device_automation_platform(
+        hass, DOMAIN, DeviceAutomationType.ACTION
+    )
+
+    with pytest.raises(InvalidDeviceAutomationConfig):
+        await platform.async_call_action_from_config(
+            hass,
+            {CONF_TYPE: "beeper.enable", CONF_DEVICE_ID: "invalid_device_id"},
+            {},
+            None,
+        )