Refactor Starlink sensor construction (home-assistant#155591)

Signed-off-by: David Rapan <[email protected]>

Author: davidrapan

homeassistant/components/starlink/sensor.py

@@ -40,21 +40,17 @@ async def async_setup_entry(
 ) -> None:
     """Set up all sensors for this entry."""
     async_add_entities(
-        StarlinkSensorEntity(config_entry.runtime_data, description)
+        description.entity_class(config_entry.runtime_data, description)
         for description in SENSORS
     )
 
-    async_add_entities(
-        StarlinkRestoreSensor(config_entry.runtime_data, description)
-        for description in RESTORE_SENSORS
-    )
-
 
 @dataclass(frozen=True, kw_only=True)
 class StarlinkSensorEntityDescription(SensorEntityDescription):
     """Describes a Starlink sensor entity."""
 
     value_fn: Callable[[StarlinkData], datetime | StateType]
+    entity_class: Callable
 
 
 class StarlinkSensorEntity(StarlinkEntity, SensorEntity):
@@ -68,14 +64,14 @@ def native_value(self) -> StateType | datetime:
         return self.entity_description.value_fn(self.coordinator.data)
 
 
-class StarlinkRestoreSensor(StarlinkSensorEntity, RestoreSensor):
-    """A RestoreSensorEntity for Starlink devices. Handles creating unique IDs."""
+class StarlinkAccumulationSensor(StarlinkSensorEntity, RestoreSensor):
+    """A StarlinkAccumulationSensor for Starlink devices. Handles creating unique IDs."""
 
     _attr_native_value: int | float = 0
 
     @property
     def native_value(self) -> int | float:
-        """Calculate the sensor value from current value and the entity description."""
+        """Calculate new value from current value and the entity description."""
         new_value = super().native_value
         if TYPE_CHECKING:
             assert isinstance(new_value, (int, float))
@@ -103,6 +99,7 @@ async def async_added_to_hass(self) -> None:
         native_unit_of_measurement=UnitOfTime.MILLISECONDS,
         suggested_display_precision=0,
         value_fn=lambda data: data.status["pop_ping_latency_ms"],
+        entity_class=StarlinkSensorEntity,
     ),
     StarlinkSensorEntityDescription(
         key="azimuth",
@@ -113,6 +110,7 @@ async def async_added_to_hass(self) -> None:
         entity_registry_enabled_default=False,
         suggested_display_precision=0,
         value_fn=lambda data: data.status["direction_azimuth"],
+        entity_class=StarlinkSensorEntity,
     ),
     StarlinkSensorEntityDescription(
         key="elevation",
@@ -123,6 +121,7 @@ async def async_added_to_hass(self) -> None:
         entity_registry_enabled_default=False,
         suggested_display_precision=0,
         value_fn=lambda data: data.status["direction_elevation"],
+        entity_class=StarlinkSensorEntity,
     ),
     StarlinkSensorEntityDescription(
         key="uplink_throughput",
@@ -133,6 +132,7 @@ async def async_added_to_hass(self) -> None:
         suggested_display_precision=1,
         suggested_unit_of_measurement=UnitOfDataRate.MEGABITS_PER_SECOND,
         value_fn=lambda data: data.status["uplink_throughput_bps"],
+        entity_class=StarlinkSensorEntity,
     ),
     StarlinkSensorEntityDescription(
         key="downlink_throughput",
@@ -143,6 +143,7 @@ async def async_added_to_hass(self) -> None:
         suggested_display_precision=1,
         suggested_unit_of_measurement=UnitOfDataRate.MEGABITS_PER_SECOND,
         value_fn=lambda data: data.status["downlink_throughput_bps"],
+        entity_class=StarlinkSensorEntity,
     ),
     StarlinkSensorEntityDescription(
         key="last_boot_time",
@@ -152,13 +153,15 @@ async def async_added_to_hass(self) -> None:
         value_fn=lambda data: (
             now() - timedelta(seconds=data.status["uptime"], milliseconds=-500)
         ).replace(microsecond=0),
+        entity_class=StarlinkSensorEntity,
     ),
     StarlinkSensorEntityDescription(
         key="ping_drop_rate",
         translation_key="ping_drop_rate",
         state_class=SensorStateClass.MEASUREMENT,
         native_unit_of_measurement=PERCENTAGE,
         value_fn=lambda data: data.status["pop_ping_drop_rate"] * 100,
+        entity_class=StarlinkSensorEntity,
     ),
     StarlinkSensorEntityDescription(
         key="power",
@@ -167,16 +170,16 @@ async def async_added_to_hass(self) -> None:
         native_unit_of_measurement=UnitOfPower.WATT,
         suggested_display_precision=0,
         value_fn=lambda data: data.consumption["latest_power"],
+        entity_class=StarlinkSensorEntity,
     ),
-)
-RESTORE_SENSORS: tuple[StarlinkSensorEntityDescription, ...] = (
     StarlinkSensorEntityDescription(
         key="energy",
         device_class=SensorDeviceClass.ENERGY,
         state_class=SensorStateClass.TOTAL_INCREASING,
         native_unit_of_measurement=UnitOfEnergy.KILO_WATT_HOUR,
         suggested_display_precision=1,
         value_fn=lambda data: data.consumption["total_energy"],
+        entity_class=StarlinkAccumulationSensor,
     ),
     StarlinkSensorEntityDescription(
         key="download",
@@ -187,6 +190,7 @@ async def async_added_to_hass(self) -> None:
         suggested_display_precision=1,
         suggested_unit_of_measurement=UnitOfInformation.GIGABYTES,
         value_fn=lambda data: data.usage["download_usage"],
+        entity_class=StarlinkAccumulationSensor,
     ),
     StarlinkSensorEntityDescription(
         key="upload",
@@ -197,5 +201,6 @@ async def async_added_to_hass(self) -> None:
         suggested_display_precision=1,
         suggested_unit_of_measurement=UnitOfInformation.GIGABYTES,
         value_fn=lambda data: data.usage["upload_usage"],
+        entity_class=StarlinkAccumulationSensor,
     ),
 )