From 8be7cb4539e6083e40fe99f6f7fad2bf27abe27a Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 8 Apr 2020 20:26:10 -0500 Subject: [PATCH] Update nut to use DataUpdateCoordinator (#33831) * Convert nut to use DataUpdateCoordinator * Adjust per review * ups_list is a dict with {id: name, ...} --- homeassistant/components/nut/__init__.py | 44 +++++++++---- homeassistant/components/nut/config_flow.py | 18 ++---- homeassistant/components/nut/const.py | 3 +- homeassistant/components/nut/sensor.py | 71 +++++++++++---------- tests/components/nut/test_config_flow.py | 6 +- 5 files changed, 81 insertions(+), 61 deletions(-) diff --git a/homeassistant/components/nut/__init__.py b/homeassistant/components/nut/__init__.py index 793dd5f2f3e..de1bcca5b31 100644 --- a/homeassistant/components/nut/__init__.py +++ b/homeassistant/components/nut/__init__.py @@ -1,7 +1,9 @@ """The nut component.""" import asyncio +from datetime import timedelta import logging +import async_timeout from pynut2.nut2 import PyNUTClient, PyNUTError from homeassistant.config_entries import ConfigEntry @@ -15,8 +17,10 @@ from homeassistant.const import ( ) from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryNotReady +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from .const import ( + COORDINATOR, DOMAIN, PLATFORMS, PYNUT_DATA, @@ -24,7 +28,6 @@ from .const import ( PYNUT_MANUFACTURER, PYNUT_MODEL, PYNUT_NAME, - PYNUT_STATUS, PYNUT_UNIQUE_ID, ) @@ -51,7 +54,22 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry): data = PyNUTData(host, port, alias, username, password) - status = await hass.async_add_executor_job(pynutdata_status, data) + async def async_update_data(): + """Fetch data from NUT.""" + async with async_timeout.timeout(10): + return await hass.async_add_executor_job(data.update) + + coordinator = DataUpdateCoordinator( + hass, + _LOGGER, + name="NUT resource status", + update_method=async_update_data, + update_interval=timedelta(seconds=60), + ) + + # Fetch initial data so we have data when entities subscribe + await coordinator.async_refresh() + status = data.status if not status: _LOGGER.error("NUT Sensor has no data, unable to set up") @@ -60,8 +78,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry): _LOGGER.debug("NUT Sensors Available: %s", status) hass.data[DOMAIN][entry.entry_id] = { + COORDINATOR: coordinator, PYNUT_DATA: data, - PYNUT_STATUS: status, PYNUT_UNIQUE_ID: _unique_id_from_status(status), PYNUT_MANUFACTURER: _manufacturer_from_status(status), PYNUT_MODEL: _model_from_status(status), @@ -143,11 +161,6 @@ def find_resources_in_config_entry(config_entry): return config_entry.data[CONF_RESOURCES] -def pynutdata_status(data): - """Wrap for data update as a callable.""" - return data.status - - async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry): """Unload a config entry.""" unload_ok = all( @@ -180,12 +193,12 @@ class PyNUTData: # Establish client with persistent=False to open/close connection on # each update call. This is more reliable with async. self._client = PyNUTClient(self._host, port, username, password, 5, False) + self.ups_list = None self._status = None @property def status(self): """Get latest update if throttle allows. Return status.""" - self.update() return self._status @property @@ -193,18 +206,21 @@ class PyNUTData: """Return the name of the ups.""" return self._alias - def list_ups(self): - """List UPSes connected to the NUT server.""" - return self._client.list_ups() - def _get_alias(self): """Get the ups alias from NUT.""" try: - return next(iter(self.list_ups())) + ups_list = self._client.list_ups() except PyNUTError as err: _LOGGER.error("Failure getting NUT ups alias, %s", err) return None + if not ups_list: + _LOGGER.error("Empty list while getting NUT ups aliases") + return None + + self.ups_list = ups_list + return list(ups_list)[0] + def _get_status(self): """Get the ups status from NUT.""" if self._alias is None: diff --git a/homeassistant/components/nut/config_flow.py b/homeassistant/components/nut/config_flow.py index 53bfe7554ea..505fd2a5e4f 100644 --- a/homeassistant/components/nut/config_flow.py +++ b/homeassistant/components/nut/config_flow.py @@ -15,7 +15,7 @@ from homeassistant.const import ( from homeassistant.core import callback import homeassistant.helpers.config_validation as cv -from . import PyNUTData, find_resources_in_config_entry, pynutdata_status +from . import PyNUTData, find_resources_in_config_entry from .const import DEFAULT_HOST, DEFAULT_PORT, SENSOR_TYPES from .const import DOMAIN # pylint:disable=unused-import @@ -54,9 +54,7 @@ def _resource_schema(available_resources, selected_resources): def _ups_schema(ups_list): """UPS selection schema.""" - ups_map = {ups: ups for ups in ups_list} - - return vol.Schema({vol.Required(CONF_ALIAS): vol.In(ups_map)}) + return vol.Schema({vol.Required(CONF_ALIAS): vol.In(ups_list)}) async def validate_input(hass: core.HomeAssistant, data): @@ -72,16 +70,12 @@ async def validate_input(hass: core.HomeAssistant, data): password = data.get(CONF_PASSWORD) data = PyNUTData(host, port, alias, username, password) - - ups_list = await hass.async_add_executor_job(data.list_ups) - if not ups_list: - raise CannotConnect - - status = await hass.async_add_executor_job(pynutdata_status, data) + await hass.async_add_executor_job(data.update) + status = data.status if not status: raise CannotConnect - return {"ups_list": ups_list, "available_resources": status} + return {"ups_list": data.ups_list, "available_resources": status} def _format_host_port_alias(user_input): @@ -135,6 +129,8 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): self.ups_list = info["ups_list"] return await self.async_step_ups() + if self._host_port_alias_already_configured(self.nut_config): + return self.async_abort(reason="already_configured") self.available_resources.update(info["available_resources"]) return await self.async_step_resources() diff --git a/homeassistant/components/nut/const.py b/homeassistant/components/nut/const.py index d7138ef865c..ae960cc4325 100644 --- a/homeassistant/components/nut/const.py +++ b/homeassistant/components/nut/const.py @@ -18,8 +18,9 @@ DEFAULT_PORT = 3493 KEY_STATUS = "ups.status" KEY_STATUS_DISPLAY = "ups.status.display" +COORDINATOR = "coordinator" + PYNUT_DATA = "data" -PYNUT_STATUS = "status" PYNUT_UNIQUE_ID = "unique_id" PYNUT_MANUFACTURER = "manufacturer" PYNUT_MODEL = "model" diff --git a/homeassistant/components/nut/sensor.py b/homeassistant/components/nut/sensor.py index 15c09001762..07abe27b426 100644 --- a/homeassistant/components/nut/sensor.py +++ b/homeassistant/components/nut/sensor.py @@ -1,5 +1,4 @@ """Provides a sensor to track various status aspects of a UPS.""" -from datetime import timedelta import logging import voluptuous as vol @@ -21,6 +20,7 @@ import homeassistant.helpers.config_validation as cv from homeassistant.helpers.entity import Entity from .const import ( + COORDINATOR, DEFAULT_HOST, DEFAULT_NAME, DEFAULT_PORT, @@ -32,7 +32,6 @@ from .const import ( PYNUT_MANUFACTURER, PYNUT_MODEL, PYNUT_NAME, - PYNUT_STATUS, PYNUT_UNIQUE_ID, SENSOR_DEVICE_CLASS, SENSOR_ICON, @@ -45,9 +44,6 @@ from .const import ( _LOGGER = logging.getLogger(__name__) -SCAN_INTERVAL = timedelta(seconds=60) - - PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend( { vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string, @@ -75,13 +71,14 @@ async def async_setup_entry(hass, config_entry, async_add_entities): """Set up the NUT sensors.""" pynut_data = hass.data[DOMAIN][config_entry.entry_id] - data = pynut_data[PYNUT_DATA] - status = pynut_data[PYNUT_STATUS] unique_id = pynut_data[PYNUT_UNIQUE_ID] manufacturer = pynut_data[PYNUT_MANUFACTURER] model = pynut_data[PYNUT_MODEL] firmware = pynut_data[PYNUT_FIRMWARE] name = pynut_data[PYNUT_NAME] + coordinator = pynut_data[COORDINATOR] + data = pynut_data[PYNUT_DATA] + status = data.status entities = [] @@ -100,8 +97,9 @@ async def async_setup_entry(hass, config_entry, async_add_entities): ): entities.append( NUTSensor( - name.title(), + coordinator, data, + name.title(), sensor_type, unique_id, manufacturer, @@ -123,10 +121,18 @@ class NUTSensor(Entity): """Representation of a sensor entity for NUT status values.""" def __init__( - self, name, data, sensor_type, unique_id, manufacturer, model, firmware + self, + coordinator, + data, + name, + sensor_type, + unique_id, + manufacturer, + model, + firmware, ): """Initialize the sensor.""" - self._data = data + self._coordinator = coordinator self._type = sensor_type self._manufacturer = manufacturer self._firmware = firmware @@ -134,10 +140,8 @@ class NUTSensor(Entity): self._device_name = name self._name = f"{name} {SENSOR_TYPES[sensor_type][SENSOR_NAME]}" self._unit = SENSOR_TYPES[sensor_type][SENSOR_UNIT] - self._state = None + self._data = data self._unique_id = unique_id - self._display_state = None - self._available = False @property def device_info(self): @@ -185,41 +189,42 @@ class NUTSensor(Entity): @property def state(self): """Return entity state from ups.""" - return self._state + if self._type == KEY_STATUS_DISPLAY: + return _format_display_state(self._data.status) + return self._data.status.get(self._type) @property def unit_of_measurement(self): """Return the unit of measurement of this entity, if any.""" return self._unit + @property + def should_poll(self): + """No need to poll. Coordinator notifies entity of updates.""" + return False + @property def available(self): - """Return if the device is polling successfully.""" - return self._available + """Return if entity is available.""" + return self._coordinator.last_update_success @property def device_state_attributes(self): """Return the sensor attributes.""" - return {ATTR_STATE: self._display_state} + return {ATTR_STATE: _format_display_state(self._data.status)} - def update(self): - """Get the latest status and use it to update our sensor state.""" - status = self._data.status + async def async_update(self): + """Update the entity. - if status is None: - self._available = False - return + Only used by the generic entity update service. + """ + await self._coordinator.async_request_refresh() - self._available = True - self._display_state = _format_display_state(status) - # In case of the display status sensor, keep a human-readable form - # as the sensor state. - if self._type == KEY_STATUS_DISPLAY: - self._state = self._display_state - elif self._type not in status: - self._state = None - else: - self._state = status[self._type] + async def async_added_to_hass(self): + """When entity is added to hass.""" + self.async_on_remove( + self._coordinator.async_add_listener(self.async_write_ha_state) + ) def _format_display_state(status): diff --git a/tests/components/nut/test_config_flow.py b/tests/components/nut/test_config_flow.py index 38953ebd235..c110df0576a 100644 --- a/tests/components/nut/test_config_flow.py +++ b/tests/components/nut/test_config_flow.py @@ -86,7 +86,8 @@ async def test_form_user_multiple_ups(hass): assert result["errors"] == {} mock_pynut = _get_mock_pynutclient( - list_vars={"battery.voltage": "voltage"}, list_ups=["ups1", "ups2"] + list_vars={"battery.voltage": "voltage"}, + list_ups={"ups1": "UPS 1", "ups2": "UPS2"}, ) with patch( @@ -146,7 +147,8 @@ async def test_form_import(hass): await setup.async_setup_component(hass, "persistent_notification", {}) mock_pynut = _get_mock_pynutclient( - list_vars={"battery.voltage": "serial"}, list_ups=["ups1"] + list_vars={"battery.voltage": "serial"}, + list_ups={"ups1": "UPS 1", "ups2": "UPS2"}, ) with patch(