Add state caching to button entities (#108272)

pull/108288/head
J. Nick Koston 2024-01-17 21:53:55 -10:00 committed by GitHub
parent 52e90b32df
commit 26cc6a5bb4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 39 additions and 16 deletions

View File

@ -1,7 +1,7 @@
"""Component to pressing a button as platforms.""" """Component to pressing a button as platforms."""
from __future__ import annotations from __future__ import annotations
from datetime import datetime, timedelta from datetime import timedelta
from enum import StrEnum from enum import StrEnum
import logging import logging
from typing import TYPE_CHECKING, final from typing import TYPE_CHECKING, final
@ -95,7 +95,7 @@ class ButtonEntity(RestoreEntity, cached_properties=CACHED_PROPERTIES_WITH_ATTR_
_attr_should_poll = False _attr_should_poll = False
_attr_device_class: ButtonDeviceClass | None _attr_device_class: ButtonDeviceClass | None
_attr_state: None = None _attr_state: None = None
__last_pressed: datetime | None = None __last_pressed_isoformat: str | None = None
def _default_to_device_class_name(self) -> bool: def _default_to_device_class_name(self) -> bool:
"""Return True if an unnamed entity should be named by its device class. """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 self.entity_description.device_class
return None return None
@property @cached_property
@final @final
def state(self) -> str | None: def state(self) -> str | None:
"""Return the entity state.""" """Return the entity state."""
if self.__last_pressed is None: return self.__last_pressed_isoformat
return None
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 @final
async def _async_press_action(self) -> None: 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. 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() self.async_write_ha_state()
await self.async_press() await self.async_press()
@ -136,7 +142,7 @@ class ButtonEntity(RestoreEntity, cached_properties=CACHED_PROPERTIES_WITH_ATTR_
await super().async_internal_added_to_hass() await super().async_internal_added_to_hass()
state = await self.async_get_last_state() state = await self.async_get_last_state()
if state is not None and state.state is not None: 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: def press(self) -> None:
"""Press the button.""" """Press the button."""

View File

@ -1,7 +1,9 @@
"""The tests for the Button component.""" """The tests for the Button component."""
from collections.abc import Generator 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 import pytest
from homeassistant.components.button import ( from homeassistant.components.button import (
@ -51,6 +53,7 @@ async def test_custom_integration(
hass: HomeAssistant, hass: HomeAssistant,
caplog: pytest.LogCaptureFixture, caplog: pytest.LogCaptureFixture,
enable_custom_integrations: None, enable_custom_integrations: None,
freezer: FrozenDateTimeFactory,
) -> None: ) -> None:
"""Test we integration.""" """Test we integration."""
platform = getattr(hass.components, f"test.{DOMAIN}") platform = getattr(hass.components, f"test.{DOMAIN}")
@ -62,7 +65,6 @@ async def test_custom_integration(
assert hass.states.get("button.button_1").state == STATE_UNKNOWN assert hass.states.get("button.button_1").state == STATE_UNKNOWN
now = dt_util.utcnow() now = dt_util.utcnow()
with patch("homeassistant.core.dt_util.utcnow", return_value=now):
await hass.services.async_call( await hass.services.async_call(
DOMAIN, DOMAIN,
SERVICE_PRESS, SERVICE_PRESS,
@ -73,6 +75,21 @@ async def test_custom_integration(
assert hass.states.get("button.button_1").state == now.isoformat() assert hass.states.get("button.button_1").state == now.isoformat()
assert "The button has been pressed" in caplog.text 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( async def test_restore_state(
hass: HomeAssistant, enable_custom_integrations: None hass: HomeAssistant, enable_custom_integrations: None