Address late comment for NextDNS (#74365)

* Simplify code

* Use async_config_entry_first_refresh()

* Use lambda to get value

homeassistant/components/nextdns/__init__.py

@@ -153,8 +153,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     except (ApiError, ClientConnectorError, asyncio.TimeoutError) as err:
         raise ConfigEntryNotReady from err
 
-    hass.data.setdefault(DOMAIN, {})
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {}
-    hass.data[DOMAIN].setdefault(entry.entry_id, {})
 
     tasks = []
 
@@ -165,7 +164,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             hass, nextdns, profile_id, update_interval
         )
         tasks.append(
-            hass.data[DOMAIN][entry.entry_id][coordinator_name].async_refresh()
+            hass.data[DOMAIN][entry.entry_id][
+                coordinator_name
+            ].async_config_entry_first_refresh()
         )
 
     await asyncio.gather(*tasks)

homeassistant/components/nextdns/sensor.py

@@ -1,8 +1,9 @@
 """Support for the NextDNS service."""
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import cast
+from typing import Any, cast
 
 from homeassistant.components.sensor import (
     SensorEntity,
@@ -14,6 +15,7 @@ from homeassistant.const import PERCENTAGE
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import NextDnsUpdateCoordinator
@@ -34,6 +36,7 @@ class NextDnsSensorRequiredKeysMixin:
     """Class for NextDNS entity required keys."""
 
     coordinator_type: str
+    value: Callable[[Any], StateType]
 
 
 @dataclass
@@ -52,6 +55,7 @@ SENSORS = (
         name="{profile_name} DNS Queries",
         native_unit_of_measurement="queries",
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.all_queries,
     ),
     NextDnsSensorEntityDescription(
         key="blocked_queries",
@@ -61,6 +65,7 @@ SENSORS = (
         name="{profile_name} DNS Queries Blocked",
         native_unit_of_measurement="queries",
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.blocked_queries,
     ),
     NextDnsSensorEntityDescription(
         key="relayed_queries",
@@ -70,6 +75,7 @@ SENSORS = (
         name="{profile_name} DNS Queries Relayed",
         native_unit_of_measurement="queries",
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.relayed_queries,
     ),
     NextDnsSensorEntityDescription(
         key="blocked_queries_ratio",
@@ -79,6 +85,7 @@ SENSORS = (
         name="{profile_name} DNS Queries Blocked Ratio",
         native_unit_of_measurement=PERCENTAGE,
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.blocked_queries_ratio,
     ),
     NextDnsSensorEntityDescription(
         key="doh_queries",
@@ -89,6 +96,7 @@ SENSORS = (
         name="{profile_name} DNS-over-HTTPS Queries",
         native_unit_of_measurement="queries",
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.doh_queries,
     ),
     NextDnsSensorEntityDescription(
         key="dot_queries",
@@ -99,6 +107,7 @@ SENSORS = (
         name="{profile_name} DNS-over-TLS Queries",
         native_unit_of_measurement="queries",
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.dot_queries,
     ),
     NextDnsSensorEntityDescription(
         key="doq_queries",
@@ -109,6 +118,7 @@ SENSORS = (
         name="{profile_name} DNS-over-QUIC Queries",
         native_unit_of_measurement="queries",
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.doq_queries,
     ),
     NextDnsSensorEntityDescription(
         key="udp_queries",
@@ -119,6 +129,7 @@ SENSORS = (
         name="{profile_name} UDP Queries",
         native_unit_of_measurement="queries",
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.udp_queries,
     ),
     NextDnsSensorEntityDescription(
         key="doh_queries_ratio",
@@ -129,6 +140,7 @@ SENSORS = (
         name="{profile_name} DNS-over-HTTPS Queries Ratio",
         native_unit_of_measurement=PERCENTAGE,
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.doh_queries_ratio,
     ),
     NextDnsSensorEntityDescription(
         key="dot_queries_ratio",
@@ -139,6 +151,7 @@ SENSORS = (
         name="{profile_name} DNS-over-TLS Queries Ratio",
         native_unit_of_measurement=PERCENTAGE,
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.dot_queries_ratio,
     ),
     NextDnsSensorEntityDescription(
         key="doq_queries_ratio",
@@ -149,6 +162,7 @@ SENSORS = (
         name="{profile_name} DNS-over-QUIC Queries Ratio",
         native_unit_of_measurement=PERCENTAGE,
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.doq_queries_ratio,
     ),
     NextDnsSensorEntityDescription(
         key="udp_queries_ratio",
@@ -159,6 +173,7 @@ SENSORS = (
         name="{profile_name} UDP Queries Ratio",
         native_unit_of_measurement=PERCENTAGE,
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.udp_queries_ratio,
     ),
     NextDnsSensorEntityDescription(
         key="encrypted_queries",
@@ -169,6 +184,7 @@ SENSORS = (
         name="{profile_name} Encrypted Queries",
         native_unit_of_measurement="queries",
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.encrypted_queries,
     ),
     NextDnsSensorEntityDescription(
         key="unencrypted_queries",
@@ -179,6 +195,7 @@ SENSORS = (
         name="{profile_name} Unencrypted Queries",
         native_unit_of_measurement="queries",
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.unencrypted_queries,
     ),
     NextDnsSensorEntityDescription(
         key="encrypted_queries_ratio",
@@ -189,6 +206,7 @@ SENSORS = (
         name="{profile_name} Encrypted Queries Ratio",
         native_unit_of_measurement=PERCENTAGE,
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.encrypted_queries_ratio,
     ),
     NextDnsSensorEntityDescription(
         key="ipv4_queries",
@@ -199,6 +217,7 @@ SENSORS = (
         name="{profile_name} IPv4 Queries",
         native_unit_of_measurement="queries",
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.ipv4_queries,
     ),
     NextDnsSensorEntityDescription(
         key="ipv6_queries",
@@ -209,6 +228,7 @@ SENSORS = (
         name="{profile_name} IPv6 Queries",
         native_unit_of_measurement="queries",
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.ipv6_queries,
     ),
     NextDnsSensorEntityDescription(
         key="ipv6_queries_ratio",
@@ -219,6 +239,7 @@ SENSORS = (
         name="{profile_name} IPv6 Queries Ratio",
         native_unit_of_measurement=PERCENTAGE,
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.ipv6_queries_ratio,
     ),
     NextDnsSensorEntityDescription(
         key="validated_queries",
@@ -229,6 +250,7 @@ SENSORS = (
         name="{profile_name} DNSSEC Validated Queries",
         native_unit_of_measurement="queries",
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.validated_queries,
     ),
     NextDnsSensorEntityDescription(
         key="not_validated_queries",
@@ -239,6 +261,7 @@ SENSORS = (
         name="{profile_name} DNSSEC Not Validated Queries",
         native_unit_of_measurement="queries",
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.not_validated_queries,
     ),
     NextDnsSensorEntityDescription(
         key="validated_queries_ratio",
@@ -249,6 +272,7 @@ SENSORS = (
         name="{profile_name} DNSSEC Validated Queries Ratio",
         native_unit_of_measurement=PERCENTAGE,
         state_class=SensorStateClass.MEASUREMENT,
+        value=lambda data: data.validated_queries_ratio,
     ),
 )
 
@@ -278,7 +302,7 @@ class NextDnsSensor(CoordinatorEntity, SensorEntity):
     def __init__(
         self,
         coordinator: NextDnsUpdateCoordinator,
-        description: SensorEntityDescription,
+        description: NextDnsSensorEntityDescription,
     ) -> None:
         """Initialize."""
         super().__init__(coordinator)
@@ -287,13 +311,11 @@ class NextDnsSensor(CoordinatorEntity, SensorEntity):
         self._attr_name = cast(str, description.name).format(
             profile_name=coordinator.profile_name
         )
-        self._attr_native_value = getattr(coordinator.data, description.key)
+        self._attr_native_value = description.value(coordinator.data)
-        self.entity_description = description
+        self.entity_description: NextDnsSensorEntityDescription = description
 
     @callback
     def _handle_coordinator_update(self) -> None:
         """Handle updated data from the coordinator."""
-        self._attr_native_value = getattr(
+        self._attr_native_value = self.entity_description.value(self.coordinator.data)
-            self.coordinator.data, self.entity_description.key
-        )
         self.async_write_ha_state()