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

@@ -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:

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()
 

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"

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):

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(