Optimistic UI updates for Vistapool write entities (#173373)

Co-authored-by: Claude <noreply@anthropic.com>

homeassistant/components/vistapool/button.py

@@ -67,7 +67,4 @@ class VistapoolLEDPulseButton(VistapoolEntity, ButtonEntity):
                 translation_key="set_failed",
                 translation_placeholders={"entity": self.entity_id},
             ) from err
-        # Optimistically reflect the just-written value so a rapid second press
-        # doesn't read the stale off-state before the Firestore push round-trips.
-        self.coordinator.data.setdefault("light", {})["status"] = 1
-        self.coordinator.async_set_updated_data(self.coordinator.data)
+        self.coordinator.apply_optimistic(_LIGHT_STATUS_PATH, 1)

homeassistant/components/vistapool/coordinator.py

@@ -81,3 +81,22 @@ class VistapoolDataUpdateCoordinator(DataUpdateCoordinator[dict[str, Any]]):
     def get_value(self, path: str, default: Any = None) -> Any:
         """Get nested data using dot-notation path."""
         return AquariteClient.get_value(self.data, path, default)
+
+    def apply_optimistic(self, value_path: str, value: Any) -> None:
+        """Reflect a just-written value before the Firestore push round-trips.
+
+        Hayward's cloud takes several seconds to acknowledge a write back
+        through Firestore, which would make the UI feel laggy. Writing into
+        coordinator.data after a successful REST call gives entities instant
+        feedback; the next snapshot from Firestore overwrites it harmlessly.
+        """
+        keys = value_path.split(".")
+        target: dict[str, Any] = self.data
+        for key in keys[:-1]:
+            child = target.get(key)
+            if not isinstance(child, dict):
+                child = {}
+                target[key] = child
+            target = child
+        target[keys[-1]] = value
+        self.async_set_updated_data(self.data)

homeassistant/components/vistapool/light.py

@@ -71,3 +71,4 @@ class VistapoolLight(VistapoolEntity, LightEntity):
                 translation_key="set_failed",
                 translation_placeholders={"entity": self.entity_id},
             ) from err
+        self.coordinator.apply_optimistic(_VALUE_PATH, value)

homeassistant/components/vistapool/number.py

@@ -233,3 +233,4 @@ class VistapoolNumber(VistapoolEntity, NumberEntity):
                 translation_key="set_failed",
                 translation_placeholders={"entity": self.entity_id},
             ) from err
+        self.coordinator.apply_optimistic(self.entity_description.value_path, raw)

tests/components/vistapool/test_init.py

@@ -103,6 +103,26 @@ async def test_setup_entry_subscribe_failure(
     assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY
 
 
+async def test_apply_optimistic_creates_missing_intermediate_dicts(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_vistapool_client: AsyncMock,
+) -> None:
+    """Test apply_optimistic walks through and creates missing intermediate dicts."""
+    mock_vistapool_client.fetch_pool_data.return_value = {"existing": "scalar"}
+    mock_config_entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    coordinator = next(iter(mock_config_entry.runtime_data.coordinators.values()))
+    coordinator.apply_optimistic("filtration.intel.temp", 27)
+    coordinator.apply_optimistic("existing.nested.key", 1)
+
+    assert coordinator.data["filtration"]["intel"]["temp"] == 27
+    assert coordinator.data["existing"] == {"nested": {"key": 1}}
+
+
 async def test_unload_entry(
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,

tests/components/vistapool/test_light.py

@@ -98,6 +98,43 @@ async def test_light_set_value(
     )
 
 
+@pytest.mark.parametrize(
+    ("service", "initial_status", "initial_state", "expected_state"),
+    [
+        pytest.param(SERVICE_TURN_ON, 0, STATE_OFF, STATE_ON, id="turn_on"),
+        pytest.param(SERVICE_TURN_OFF, 1, STATE_ON, STATE_OFF, id="turn_off"),
+    ],
+)
+async def test_light_optimistic_state(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_vistapool_client: AsyncMock,
+    mock_pool_data: dict[str, Any],
+    service: str,
+    initial_status: int,
+    initial_state: str,
+    expected_state: str,
+) -> None:
+    """Test the entity state reflects the just-written value before the Firestore push."""
+    mock_pool_data["light"] = {"status": initial_status}
+    mock_vistapool_client.fetch_pool_data.return_value = mock_pool_data
+    mock_config_entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("light.my_pool_light").state == initial_state
+
+    await hass.services.async_call(
+        LIGHT_DOMAIN,
+        service,
+        {ATTR_ENTITY_ID: "light.my_pool_light"},
+        blocking=True,
+    )
+
+    assert hass.states.get("light.my_pool_light").state == expected_state
+
+
 async def test_light_set_value_raises_on_api_error(
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,

tests/components/vistapool/test_number.py

@@ -232,6 +232,7 @@ async def test_number_set_value(
     mock_vistapool_client.set_value.assert_awaited_once_with(
         "ABCDEF1234567890", expected_path, expected_raw
     )
+    assert hass.states.get(entity_id).state == str(float(user_value))
     value_arg = mock_vistapool_client.set_value.await_args.args[2]
     assert isinstance(value_arg, int)