Starlink's Energy, Download and Upload accumulation after restart fix (#137855)

Co-authored-by: Erik Montnemery <erik@montnemery.com>

homeassistant/components/starlink/coordinator.py

@@ -66,6 +66,7 @@ class StarlinkUpdateCoordinator(DataUpdateCoordinator[StarlinkData]):
             config_entry=config_entry,
             name=config_entry.title,
             update_interval=timedelta(seconds=5),
+            always_update=False,
         )
 
     def _get_starlink_data(self) -> StarlinkData:
@@ -76,17 +77,11 @@ class StarlinkUpdateCoordinator(DataUpdateCoordinator[StarlinkData]):
         sleep = get_sleep_config(context)
         status, obstruction, alert = status_data(context)
         index, _, _, _, _, usage, consumption, *_ = history_stats(
-            parse_samples=-1, start=self.history_stats_start, context=context
+            parse_samples=-1 if self.history_stats_start is not None else 1,
+            start=self.history_stats_start,
+            context=context,
         )
         self.history_stats_start = index["end_counter"]
-        if self.data:
-            if index["samples"] > 0:
-                usage["download_usage"] += self.data.usage["download_usage"]
-                usage["upload_usage"] += self.data.usage["upload_usage"]
-                consumption["total_energy"] += self.data.consumption["total_energy"]
-            else:
-                usage = self.data.usage
-                consumption = self.data.consumption
         return StarlinkData(
             location, sleep, status, obstruction, alert, usage, consumption
         )
@@ -94,10 +89,9 @@ class StarlinkUpdateCoordinator(DataUpdateCoordinator[StarlinkData]):
     async def _async_update_data(self) -> StarlinkData:
         async with asyncio.timeout(4):
             try:
-                result = await self.hass.async_add_executor_job(self._get_starlink_data)
+                return await self.hass.async_add_executor_job(self._get_starlink_data)
             except GrpcError as exc:
                 raise UpdateFailed from exc
-            return result
 
     async def async_stow_starlink(self, stow: bool) -> None:
         """Set whether Starlink system tied to this coordinator should be stowed."""

homeassistant/components/starlink/sensor.py

@@ -5,8 +5,10 @@ from __future__ import annotations
 from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import datetime, timedelta
+from typing import TYPE_CHECKING
 
 from homeassistant.components.sensor import (
+    RestoreSensor,
     SensorDeviceClass,
     SensorEntity,
     SensorEntityDescription,
@@ -42,6 +44,11 @@ async def async_setup_entry(
         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):
@@ -61,6 +68,33 @@ class StarlinkSensorEntity(StarlinkEntity, SensorEntity):
         return self.entity_description.value_fn(self.coordinator.data)
 
 
+class StarlinkRestoreSensor(StarlinkSensorEntity, RestoreSensor):
+    """A RestoreSensorEntity 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."""
+        new_value = super().native_value
+        if TYPE_CHECKING:
+            assert isinstance(new_value, (int, float))
+        self._attr_native_value += new_value
+        return self._attr_native_value
+
+    async def async_added_to_hass(self) -> None:
+        """When entity is added to hass."""
+        await super().async_added_to_hass()
+        if (
+            last_sensor_data := await self.async_get_last_sensor_data()
+        ) is not None and (
+            last_native_value := last_sensor_data.native_value
+        ) is not None:
+            if TYPE_CHECKING:
+                assert isinstance(last_native_value, (int, float))
+            self._attr_native_value = last_native_value
+
+
 SENSORS: tuple[StarlinkSensorEntityDescription, ...] = (
     StarlinkSensorEntityDescription(
         key="ping",
@@ -96,7 +130,8 @@ SENSORS: tuple[StarlinkSensorEntityDescription, ...] = (
         state_class=SensorStateClass.MEASUREMENT,
         device_class=SensorDeviceClass.DATA_RATE,
         native_unit_of_measurement=UnitOfDataRate.BITS_PER_SECOND,
-        suggested_display_precision=0,
+        suggested_display_precision=1,
+        suggested_unit_of_measurement=UnitOfDataRate.MEGABITS_PER_SECOND,
         value_fn=lambda data: data.status["uplink_throughput_bps"],
     ),
     StarlinkSensorEntityDescription(
@@ -105,7 +140,8 @@ SENSORS: tuple[StarlinkSensorEntityDescription, ...] = (
         state_class=SensorStateClass.MEASUREMENT,
         device_class=SensorDeviceClass.DATA_RATE,
         native_unit_of_measurement=UnitOfDataRate.BITS_PER_SECOND,
-        suggested_display_precision=0,
+        suggested_display_precision=1,
+        suggested_unit_of_measurement=UnitOfDataRate.MEGABITS_PER_SECOND,
         value_fn=lambda data: data.status["downlink_throughput_bps"],
     ),
     StarlinkSensorEntityDescription(
@@ -125,13 +161,22 @@ SENSORS: tuple[StarlinkSensorEntityDescription, ...] = (
         value_fn=lambda data: data.status["pop_ping_drop_rate"] * 100,
     ),
     StarlinkSensorEntityDescription(
-        key="upload",
-        translation_key="upload",
-        device_class=SensorDeviceClass.DATA_SIZE,
+        key="power",
+        device_class=SensorDeviceClass.POWER,
+        state_class=SensorStateClass.MEASUREMENT,
+        native_unit_of_measurement=UnitOfPower.WATT,
+        suggested_display_precision=0,
+        value_fn=lambda data: data.consumption["latest_power"],
+    ),
+)
+RESTORE_SENSORS: tuple[StarlinkSensorEntityDescription, ...] = (
+    StarlinkSensorEntityDescription(
+        key="energy",
+        device_class=SensorDeviceClass.ENERGY,
         state_class=SensorStateClass.TOTAL_INCREASING,
-        native_unit_of_measurement=UnitOfInformation.BYTES,
-        suggested_unit_of_measurement=UnitOfInformation.GIGABYTES,
-        value_fn=lambda data: data.usage["upload_usage"],
+        native_unit_of_measurement=UnitOfEnergy.KILO_WATT_HOUR,
+        suggested_display_precision=1,
+        value_fn=lambda data: data.consumption["total_energy"],
     ),
     StarlinkSensorEntityDescription(
         key="download",
@@ -139,21 +184,18 @@ SENSORS: tuple[StarlinkSensorEntityDescription, ...] = (
         device_class=SensorDeviceClass.DATA_SIZE,
         state_class=SensorStateClass.TOTAL_INCREASING,
         native_unit_of_measurement=UnitOfInformation.BYTES,
+        suggested_display_precision=1,
         suggested_unit_of_measurement=UnitOfInformation.GIGABYTES,
         value_fn=lambda data: data.usage["download_usage"],
     ),
     StarlinkSensorEntityDescription(
-        key="power",
-        device_class=SensorDeviceClass.POWER,
-        state_class=SensorStateClass.MEASUREMENT,
-        native_unit_of_measurement=UnitOfPower.WATT,
-        value_fn=lambda data: data.consumption["latest_power"],
-    ),
-    StarlinkSensorEntityDescription(
-        key="energy",
-        device_class=SensorDeviceClass.ENERGY,
+        key="upload",
+        translation_key="upload",
+        device_class=SensorDeviceClass.DATA_SIZE,
         state_class=SensorStateClass.TOTAL_INCREASING,
-        native_unit_of_measurement=UnitOfEnergy.KILO_WATT_HOUR,
-        value_fn=lambda data: data.consumption["total_energy"],
+        native_unit_of_measurement=UnitOfInformation.BYTES,
+        suggested_display_precision=1,
+        suggested_unit_of_measurement=UnitOfInformation.GIGABYTES,
+        value_fn=lambda data: data.usage["upload_usage"],
     ),
 )

tests/components/starlink/test_init.py

@@ -1,9 +1,11 @@
 """Tests Starlink integration init/unload."""
 
+from unittest.mock import patch
+
 from homeassistant.components.starlink.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import CONF_IP_ADDRESS
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, State
 
 from .patchers import (
     HISTORY_STATS_SUCCESS_PATCHER,
@@ -12,7 +14,7 @@ from .patchers import (
     STATUS_DATA_SUCCESS_PATCHER,
 )
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, mock_restore_cache_with_extra_data
 
 
 async def test_successful_entry(hass: HomeAssistant) -> None:
@@ -60,3 +62,53 @@ async def test_unload_entry(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
 
         assert entry.state is ConfigEntryState.NOT_LOADED
+
+
+async def test_restore_cache_with_accumulation(hass: HomeAssistant) -> None:
+    """Test configuring Starlink."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={CONF_IP_ADDRESS: "1.2.3.4:0000"},
+    )
+    entity_id = "sensor.starlink_energy"
+
+    mock_restore_cache_with_extra_data(
+        hass,
+        (
+            (
+                State(
+                    entity_id,
+                    "",
+                ),
+                {
+                    "native_value": 1,
+                    "native_unit_of_measurement": None,
+                },
+            ),
+        ),
+    )
+
+    with (
+        STATUS_DATA_SUCCESS_PATCHER,
+        LOCATION_DATA_SUCCESS_PATCHER,
+        SLEEP_DATA_SUCCESS_PATCHER,
+        HISTORY_STATS_SUCCESS_PATCHER,
+    ):
+        entry.add_to_hass(hass)
+
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+        assert entry.runtime_data
+        assert entry.runtime_data.data
+
+        assert hass.states.get(entity_id).state == str(1 + 0.00786231368489)
+
+        await entry.runtime_data.async_refresh()
+
+        assert hass.states.get(entity_id).state == str(1 + 0.00786231368489)
+
+        with patch.object(entry.runtime_data, "always_update", return_value=True):
+            await entry.runtime_data.async_refresh()
+
+        assert hass.states.get(entity_id).state == str(1 + 0.01572462736977)