Add data update coordinator to Whois (#64846)

Co-authored-by: Joakim Sørensen <joasoe@gmail.com>

homeassistant/components/whois/__init__.py

@@ -1,13 +1,48 @@
 """The Whois integration."""
-from homeassistant.config_entries import ConfigEntry
+from __future__ import annotations
-from homeassistant.core import HomeAssistant
 
-from .const import PLATFORMS
+from whois import Domain, query as whois_query
+from whois.exceptions import (
+    FailedParsingWhoisOutput,
+    UnknownDateFormat,
+    UnknownTld,
+    WhoisCommandFailed,
+)
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_DOMAIN
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DOMAIN, LOGGER, PLATFORMS, SCAN_INTERVAL
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up from a config entry."""
+
+    async def _async_query_domain() -> Domain | None:
+        """Query WHOIS for domain information."""
+        try:
+            return await hass.async_add_executor_job(
+                whois_query, entry.data[CONF_DOMAIN]
+            )
+        except UnknownTld as ex:
+            raise UpdateFailed("Could not set up whois, TLD is unknown") from ex
+        except (FailedParsingWhoisOutput, WhoisCommandFailed, UnknownDateFormat) as ex:
+            raise UpdateFailed("An error occurred during WHOIS lookup") from ex
+
+    coordinator: DataUpdateCoordinator[Domain | None] = DataUpdateCoordinator(
+        hass,
+        LOGGER,
+        name=f"{DOMAIN}_APK",
+        update_interval=SCAN_INTERVAL,
+        update_method=_async_query_domain,
+    )
+    await coordinator.async_config_entry_first_refresh()
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
+
     return True
 
 

homeassistant/components/whois/const.py

@@ -1,6 +1,7 @@
 """Constants for the Whois integration."""
 from __future__ import annotations
 
+from datetime import timedelta
 import logging
 from typing import Final
 
@@ -11,6 +12,8 @@ PLATFORMS = [Platform.SENSOR]
 
 LOGGER = logging.getLogger(__package__)
 
+SCAN_INTERVAL = timedelta(hours=24)
+
 DEFAULT_NAME = "Whois"
 
 ATTR_EXPIRES = "expires"

homeassistant/components/whois/sensor.py

@@ -3,18 +3,10 @@ from __future__ import annotations
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from datetime import timedelta
 from typing import cast
 
 import voluptuous as vol
-import whois
 from whois import Domain
-from whois.exceptions import (
-    FailedParsingWhoisOutput,
-    UnknownDateFormat,
-    UnknownTld,
-    WhoisCommandFailed,
-)
 
 from homeassistant.components.sensor import (
     PLATFORM_SCHEMA,
@@ -29,6 +21,10 @@ from homeassistant.helpers.device_registry import DeviceEntryType
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
+from homeassistant.helpers.update_coordinator import (
+    CoordinatorEntity,
+    DataUpdateCoordinator,
+)
 
 from .const import (
     ATTR_EXPIRES,
@@ -40,8 +36,6 @@ from .const import (
     LOGGER,
 )
 
-SCAN_INTERVAL = timedelta(hours=24)
-
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
     {
         vol.Required(CONF_DOMAIN): cv.string,
@@ -111,21 +105,13 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the platform from config_entry."""
-    domain = entry.data[CONF_DOMAIN]
+    coordinator = hass.data[DOMAIN][entry.entry_id]
-    try:
-        await hass.async_add_executor_job(whois.query, domain)
-    except UnknownTld:
-        LOGGER.error("Could not set up whois for %s, TLD is unknown", domain)
-        return
-    except (FailedParsingWhoisOutput, WhoisCommandFailed, UnknownDateFormat) as ex:
-        LOGGER.error("Exception %s occurred during WHOIS lookup for %s", ex, domain)
-        return
-
     async_add_entities(
         [
             WhoisSensorEntity(
-                domain=domain,
+                coordinator=coordinator,
                 description=description,
+                domain=entry.data[CONF_DOMAIN],
             )
             for description in SENSORS
         ],
@@ -133,13 +119,19 @@ async def async_setup_entry(
     )
 
 
-class WhoisSensorEntity(SensorEntity):
+class WhoisSensorEntity(CoordinatorEntity, SensorEntity):
     """Implementation of a WHOIS sensor."""
 
     entity_description: WhoisSensorEntityDescription
 
-    def __init__(self, description: WhoisSensorEntityDescription, domain: str) -> None:
+    def __init__(
+        self,
+        coordinator: DataUpdateCoordinator,
+        description: WhoisSensorEntityDescription,
+        domain: str,
+    ) -> None:
         """Initialize the sensor."""
+        super().__init__(coordinator=coordinator)
         self.entity_description = description
         self._attr_name = domain
         self._attr_unique_id = f"{domain}_{description.key}"
@@ -149,42 +141,35 @@ class WhoisSensorEntity(SensorEntity):
         )
         self._domain = domain
 
-    def _empty_value_and_attributes(self) -> None:
+    @property
-        """Empty the state and attributes on an error."""
+    def native_value(self) -> int | None:
-        self._attr_native_value = None
+        """Return the state of the sensor."""
-        self._attr_extra_state_attributes = {}
+        if self.coordinator.data is None:
+            return None
+        return self.entity_description.value_fn(self.coordinator.data)
 
-    def update(self) -> None:
+    @property
-        """Get the current WHOIS data for the domain."""
+    def extra_state_attributes(self) -> dict[str, int | float | None] | None:
-        try:
+        """Return the state attributes of the monitored installation."""
-            response: Domain | None = whois.query(self._domain)
-        except (FailedParsingWhoisOutput, WhoisCommandFailed, UnknownDateFormat) as ex:
-            LOGGER.error("Exception %s occurred during WHOIS lookup", ex)
-            self._empty_value_and_attributes()
-            return
 
-        if response:
+        # Only add attributes to the original sensor
-            if not response.expiration_date:
+        if self.entity_description.key != "days_until_expiration":
-                LOGGER.error("Failed to find expiration_date in whois lookup response")
+            return None
-                self._empty_value_and_attributes()
-                return
 
-            self._attr_native_value = self.entity_description.value_fn(response)
+        if self.coordinator.data is None:
+            return None
 
-            # Only add attributes to the original sensor
+        attrs = {
-            if self.entity_description.key != "days_until_expiration":
+            ATTR_EXPIRES: self.coordinator.data.expiration_date.isoformat(),
-                return None
+        }
 
-            attrs = {}
+        if self.coordinator.data.name_servers:
-            attrs[ATTR_EXPIRES] = response.expiration_date.isoformat()
+            attrs[ATTR_NAME_SERVERS] = " ".join(self.coordinator.data.name_servers)
 
-            if response.name_servers:
+        if self.coordinator.data.last_updated:
-                attrs[ATTR_NAME_SERVERS] = " ".join(response.name_servers)
+            attrs[ATTR_UPDATED] = self.coordinator.data.last_updated.isoformat()
 
-            if response.last_updated:
+        if self.coordinator.data.registrar:
-                attrs[ATTR_UPDATED] = response.last_updated.isoformat()
+            attrs[ATTR_REGISTRAR] = self.coordinator.data.registrar
 
-            if response.registrar:
+        return attrs
-                attrs[ATTR_REGISTRAR] = response.registrar
-
-            self._attr_extra_state_attributes = attrs