Keep Portainer entities known on endpoint timeout (#184325)

homeassistant/components/portainer/coordinator.py

@@ -224,6 +224,8 @@ class PortainerCoordinator(
         self._event_listeners: dict[int, PortainerEventListener] = {}
         self._event_listeners_enabled = False
         self._container_ids_by_endpoint: dict[int, dict[str, str]] = {}
+        # Last successful fetch per endpoint, kept across timed-out polls
+        self._last_endpoint_data: dict[int, PortainerCoordinatorData] = {}
 
     @override
     async def update_data(self) -> dict[int, PortainerCoordinatorData]:
@@ -246,6 +248,7 @@ class PortainerCoordinator(
             ) from err
 
         mapped_endpoints: dict[int, PortainerCoordinatorData] = {}
+        timed_out_endpoints: set[int] = set()
         for endpoint in endpoints:
             if endpoint.status == EndpointStatus.DOWN:
                 _LOGGER.debug(
@@ -293,7 +296,7 @@ class PortainerCoordinator(
                     for stack in result
                 ]
 
-                prev_endpoint = self.data.get(endpoint.id) if self.data else None
+                prev_endpoint = self._last_endpoint_data.get(endpoint.id)
                 container_map: dict[str, PortainerContainerData] = {}
                 stack_map: dict[str, PortainerStackData] = {
                     stack.name: PortainerStackData(stack=stack, container_count=0)
@@ -433,9 +436,15 @@ class PortainerCoordinator(
                     endpoint.name,
                     endpoint.id,
                 )
+                timed_out_endpoints.add(endpoint.id)
                 continue
 
-        self._async_add_remove_endpoints(mapped_endpoints)
+        self._last_endpoint_data = {
+            endpoint_id: data
+            for endpoint_id, data in self._last_endpoint_data.items()
+            if endpoint_id in timed_out_endpoints
+        } | mapped_endpoints
+        self._async_add_remove_endpoints(mapped_endpoints, timed_out_endpoints)
         self._async_sync_event_listeners(mapped_endpoints)
 
         return mapped_endpoints
@@ -492,11 +501,17 @@ class PortainerCoordinator(
             )
 
     def _async_add_remove_endpoints(
-        self, mapped_endpoints: dict[int, PortainerCoordinatorData]
+        self,
+        mapped_endpoints: dict[int, PortainerCoordinatorData],
+        timed_out_endpoints: set[int],
     ) -> None:
-        """Add new endpoints, remove non-existing endpoints."""
+        """Add new endpoints, remove non-existing endpoints.
+
+        Timed-out endpoints keep their known entities, so they aren't added
+        again once the endpoint answers.
+        """
         current_endpoints = {endpoint.id for endpoint in mapped_endpoints.values()}
-        self.known_endpoints &= current_endpoints
+        self.known_endpoints &= current_endpoints | timed_out_endpoints
         new_endpoints = current_endpoints - self.known_endpoints
 
         # The stack ID is part of the key because it is part of the stack device
@@ -507,7 +522,11 @@ class PortainerCoordinator(
             for endpoint in mapped_endpoints.values()
             for stack_name, stack_data in endpoint.stacks.items()
         }
-        self.known_stacks &= current_stacks
+        self.known_stacks = {
+            stack
+            for stack in self.known_stacks
+            if stack in current_stacks or stack[0] in timed_out_endpoints
+        }
         new_stacks = current_stacks - self.known_stacks
 
         if new_endpoints or new_stacks:
@@ -534,7 +553,11 @@ class PortainerCoordinator(
         }
         # Prune departed containers so a recreated container is detected as new
         # and its entity is rebuilt with the fresh (ephemeral) Docker container ID.
-        self.known_containers &= current_containers
+        self.known_containers = {
+            container
+            for container in self.known_containers
+            if container in current_containers or container[0] in timed_out_endpoints
+        }
         new_containers = current_containers - self.known_containers
         if new_containers:
             _LOGGER.debug("New containers found: %s", new_containers)
@@ -556,7 +579,11 @@ class PortainerCoordinator(
             for volume_name in endpoint.volumes
         }
 
-        self.known_volumes &= current_volumes
+        self.known_volumes = {
+            volume
+            for volume in self.known_volumes
+            if volume in current_volumes or volume[0] in timed_out_endpoints
+        }
         new_volumes = current_volumes - self.known_volumes
         if new_volumes:
             _LOGGER.debug("New volumes found: %s", new_volumes)

tests/components/portainer/test_sensor.py

@@ -5,14 +5,17 @@ from unittest.mock import AsyncMock, patch
 
 from freezegun.api import FrozenDateTimeFactory
 from pyportainer.exceptions import PortainerTimeoutError
-from pyportainer.models.docker import EndpointStatus
+from pyportainer.models.docker import DockerContainerStats, EndpointStatus
 from pyportainer.models.docker_inspect import DockerInspect
 from pyportainer.models.portainer import Endpoint
 import pytest
 from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.components.portainer.const import DOMAIN
-from homeassistant.components.portainer.coordinator import DEFAULT_DF_SCAN_INTERVAL
+from homeassistant.components.portainer.coordinator import (
+    DEFAULT_DF_SCAN_INTERVAL,
+    DEFAULT_SCAN_INTERVAL,
+)
 from homeassistant.const import STATE_UNAVAILABLE, STATE_UNKNOWN, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
@@ -28,6 +31,8 @@ from tests.common import (
     snapshot_platform,
 )
 
+CPU_USAGE_ENTITY_ID = "sensor.focused_einstein_cpu_usage_total"
+
 
 @pytest.fixture(autouse=True)
 def enable_all_entities(entity_registry_enabled_by_default: None) -> None:
@@ -105,6 +110,43 @@ async def test_df_endpoint_timeout_only_marks_that_endpoint_unavailable(
     assert state.state != STATE_UNAVAILABLE
 
 
+async def test_endpoint_timeout_keeps_known_entities(
+    hass: HomeAssistant,
+    mock_portainer_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    freezer: FrozenDateTimeFactory,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test a timed-out poll doesn't re-add entities or reset CPU usage."""
+    await setup_integration(hass, mock_config_entry)
+
+    mock_portainer_client.docker_version.side_effect = PortainerTimeoutError("timeout")
+    freezer.tick(DEFAULT_SCAN_INTERVAL)
+    async_fire_time_changed(hass, dt_util.utcnow())
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    assert (state := hass.states.get(CPU_USAGE_ENTITY_ID))
+    assert state.state == STATE_UNAVAILABLE
+
+    stats = cast(
+        dict[str, Any], load_json_value_fixture("container_stats.json", DOMAIN)
+    )
+    stats["cpu_stats"]["system_cpu_usage"] += 1_000_000_000
+    stats["cpu_stats"]["cpu_usage"]["total_usage"] += 125_000_000
+    mock_portainer_client.container_stats.return_value = DockerContainerStats.from_dict(
+        stats
+    )
+    mock_portainer_client.docker_version.side_effect = None
+    freezer.tick(DEFAULT_SCAN_INTERVAL)
+    async_fire_time_changed(hass, dt_util.utcnow())
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    # Usage is averaged since the last successful poll
+    assert (state := hass.states.get(CPU_USAGE_ENTITY_ID))
+    assert state.state == "50.0"
+    assert "does not generate unique IDs" not in caplog.text
+
+
 async def test_container_never_started(
     hass: HomeAssistant,
     mock_portainer_client: AsyncMock,