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."""
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."""

View File

@ -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