From 90a3c2e35748c97bdcad664deda56069feb49ef6 Mon Sep 17 00:00:00 2001 From: Michael <35783820+mib1185@users.noreply.github.com> Date: Sat, 4 May 2024 23:08:01 +0200 Subject: [PATCH] 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 | 32 +++++++++-------- homeassistant/components/nut/const.py | 7 ---- homeassistant/components/nut/device_action.py | 36 ++++++++++--------- homeassistant/components/nut/diagnostics.py | 15 ++++---- homeassistant/components/nut/sensor.py | 23 ++++-------- tests/components/nut/test_device_action.py | 28 ++++++++++++++- 6 files changed, 77 insertions(+), 64 deletions(-) diff --git a/homeassistant/components/nut/__init__.py b/homeassistant/components/nut/__init__.py index 8b715237e01..640dbb1416a 100644 --- a/homeassistant/components/nut/__init__.py +++ b/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: diff --git a/homeassistant/components/nut/const.py b/homeassistant/components/nut/const.py index 9be06de1f73..6db40a910a0 100644 --- a/homeassistant/components/nut/const.py +++ b/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", diff --git a/homeassistant/components/nut/device_action.py b/homeassistant/components/nut/device_action.py index 0ec58e651b2..a051f843226 100644 --- a/homeassistant/components/nut/device_action.py +++ b/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 diff --git a/homeassistant/components/nut/diagnostics.py b/homeassistant/components/nut/diagnostics.py index 88a05e461c9..532e4ece76b 100644 --- a/homeassistant/components/nut/diagnostics.py +++ b/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 diff --git a/homeassistant/components/nut/sensor.py b/homeassistant/components/nut/sensor.py index cd5ae64901d..7b61342866b 100644 --- a/homeassistant/components/nut/sensor.py +++ b/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] diff --git a/tests/components/nut/test_device_action.py b/tests/components/nut/test_device_action.py index 8113b19e313..01675f928e3 100644 --- a/tests/components/nut/test_device_action.py +++ b/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, + )