From 26cc6a5bb4383217d241f18e23241aebfc59cd31 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Wed, 17 Jan 2024 21:53:55 -1000
Subject: [PATCH] Add state caching to button entities (#108272)

---
 homeassistant/components/button/__init__.py | 22 +++++++++-----
 tests/components/button/test_init.py        | 33 ++++++++++++++++-----
 2 files changed, 39 insertions(+), 16 deletions(-)

diff --git a/homeassistant/components/button/__init__.py b/homeassistant/components/button/__init__.py
index 358348a8077..3ecc27f8573 100644
--- a/homeassistant/components/button/__init__.py
+++ b/homeassistant/components/button/__init__.py
@@ -1,7 +1,7 @@
 """Component to pressing a button as platforms."""
 from __future__ import annotations
 
-from datetime import datetime, timedelta
+from datetime import timedelta
 from enum import StrEnum
 import logging
 from typing import TYPE_CHECKING, final
@@ -95,7 +95,7 @@ class ButtonEntity(RestoreEntity, cached_properties=CACHED_PROPERTIES_WITH_ATTR_
     _attr_should_poll = False
     _attr_device_class: ButtonDeviceClass | None
     _attr_state: None = None
-    __last_pressed: datetime | None = None
+    __last_pressed_isoformat: str | None = None
 
     def _default_to_device_class_name(self) -> bool:
         """Return True if an unnamed entity should be named by its device class.
@@ -113,13 +113,19 @@ class ButtonEntity(RestoreEntity, cached_properties=CACHED_PROPERTIES_WITH_ATTR_
             return self.entity_description.device_class
         return None
 
-    @property
+    @cached_property
     @final
     def state(self) -> str | None:
         """Return the entity state."""
-        if self.__last_pressed is None:
-            return None
-        return self.__last_pressed.isoformat()
+        return self.__last_pressed_isoformat
+
+    def __set_state(self, state: str | None) -> None:
+        """Set the entity state."""
+        try:  # noqa: SIM105  suppress is much slower
+            del self.state
+        except AttributeError:
+            pass
+        self.__last_pressed_isoformat = state
 
     @final
     async def _async_press_action(self) -> None:
@@ -127,7 +133,7 @@ class ButtonEntity(RestoreEntity, cached_properties=CACHED_PROPERTIES_WITH_ATTR_
 
         Should not be overridden, handle setting last press timestamp.
         """
-        self.__last_pressed = dt_util.utcnow()
+        self.__set_state(dt_util.utcnow().isoformat())
         self.async_write_ha_state()
         await self.async_press()
 
@@ -136,7 +142,7 @@ class ButtonEntity(RestoreEntity, cached_properties=CACHED_PROPERTIES_WITH_ATTR_
         await super().async_internal_added_to_hass()
         state = await self.async_get_last_state()
         if state is not None and state.state is not None:
-            self.__last_pressed = dt_util.parse_datetime(state.state)
+            self.__set_state(state.state)
 
     def press(self) -> None:
         """Press the button."""
diff --git a/tests/components/button/test_init.py b/tests/components/button/test_init.py
index 24f893578ce..2457a796d45 100644
--- a/tests/components/button/test_init.py
+++ b/tests/components/button/test_init.py
@@ -1,7 +1,9 @@
 """The tests for the Button component."""
 from collections.abc import Generator
-from unittest.mock import MagicMock, patch
+from datetime import timedelta
+from unittest.mock import MagicMock
 
+from freezegun.api import FrozenDateTimeFactory
 import pytest
 
 from homeassistant.components.button import (
@@ -51,6 +53,7 @@ async def test_custom_integration(
     hass: HomeAssistant,
     caplog: pytest.LogCaptureFixture,
     enable_custom_integrations: None,
+    freezer: FrozenDateTimeFactory,
 ) -> None:
     """Test we integration."""
     platform = getattr(hass.components, f"test.{DOMAIN}")
@@ -62,17 +65,31 @@ async def test_custom_integration(
     assert hass.states.get("button.button_1").state == STATE_UNKNOWN
 
     now = dt_util.utcnow()
-    with patch("homeassistant.core.dt_util.utcnow", return_value=now):
-        await hass.services.async_call(
-            DOMAIN,
-            SERVICE_PRESS,
-            {ATTR_ENTITY_ID: "button.button_1"},
-            blocking=True,
-        )
+    await hass.services.async_call(
+        DOMAIN,
+        SERVICE_PRESS,
+        {ATTR_ENTITY_ID: "button.button_1"},
+        blocking=True,
+    )
 
     assert hass.states.get("button.button_1").state == now.isoformat()
     assert "The button has been pressed" in caplog.text
 
+    now_isoformat = dt_util.utcnow().isoformat()
+    assert hass.states.get("button.button_1").state == now_isoformat
+
+    new_time = dt_util.utcnow() + timedelta(weeks=1)
+    freezer.move_to(new_time)
+    await hass.services.async_call(
+        DOMAIN,
+        SERVICE_PRESS,
+        {ATTR_ENTITY_ID: "button.button_1"},
+        blocking=True,
+    )
+
+    new_time_isoformat = new_time.isoformat()
+    assert hass.states.get("button.button_1").state == new_time_isoformat
+
 
 async def test_restore_state(
     hass: HomeAssistant, enable_custom_integrations: None