Fix ESPHome button not getting device updates (#95311)

.coveragerc

@@ -308,7 +308,6 @@ omit =
     homeassistant/components/escea/discovery.py
     homeassistant/components/esphome/__init__.py
     homeassistant/components/esphome/bluetooth/*
-    homeassistant/components/esphome/button.py
     homeassistant/components/esphome/camera.py
     homeassistant/components/esphome/cover.py
     homeassistant/components/esphome/domain_data.py

homeassistant/components/esphome/button.py

@@ -42,10 +42,18 @@ class EsphomeButton(EsphomeEntity[ButtonInfo, EntityState], ButtonEntity):
 
     @callback
     def _on_device_update(self) -> None:
-        """Update the entity state when device info has changed."""
-        # This override the EsphomeEntity method as the button entity
-        # never gets a state update.
-        self._on_state_update()
+        """Call when device updates or entry data changes.
+
+        The default behavior is only to write entity state when the
+        device is unavailable when the device state changes.
+        This method overrides the default behavior since buttons do
+        not have a state, so we will never get a state update for a
+        button. As such, we need to write the state on every device
+        update to ensure the button goes available and unavailable
+        as the device becomes available or unavailable.
+        """
+        self._on_entry_data_changed()
+        self.async_write_ha_state()
 
     async def async_press(self) -> None:
         """Press the button."""

tests/components/esphome/conftest.py

@@ -153,6 +153,7 @@ class MockESPHomeDevice:
         """Init the mock."""
         self.entry = entry
         self.state_callback: Callable[[EntityState], None]
+        self.on_disconnect: Callable[[bool], None]
 
     def set_state_callback(self, state_callback: Callable[[EntityState], None]) -> None:
         """Set the state callback."""
@@ -162,6 +163,14 @@ class MockESPHomeDevice:
         """Mock setting state."""
         self.state_callback(state)
 
+    def set_on_disconnect(self, on_disconnect: Callable[[bool], None]) -> None:
+        """Set the disconnect callback."""
+        self.on_disconnect = on_disconnect
+
+    async def mock_disconnect(self, expected_disconnect: bool) -> None:
+        """Mock disconnecting."""
+        await self.on_disconnect(expected_disconnect)
+
 
 async def _mock_generic_device_entry(
     hass: HomeAssistant,
@@ -209,15 +218,23 @@ async def _mock_generic_device_entry(
     mock_client.subscribe_states = _subscribe_states
 
     try_connect_done = Event()
-    real_try_connect = ReconnectLogic._try_connect
 
-    async def mock_try_connect(self):
-        """Set an event when ReconnectLogic._try_connect has been awaited."""
-        result = await real_try_connect(self)
-        try_connect_done.set()
-        return result
+    class MockReconnectLogic(ReconnectLogic):
+        """Mock ReconnectLogic."""
 
-    with patch.object(ReconnectLogic, "_try_connect", mock_try_connect):
+        def __init__(self, *args, **kwargs):
+            """Init the mock."""
+            super().__init__(*args, **kwargs)
+            mock_device.set_on_disconnect(kwargs["on_disconnect"])
+            self._try_connect = self.mock_try_connect
+
+        async def mock_try_connect(self):
+            """Set an event when ReconnectLogic._try_connect has been awaited."""
+            result = await super()._try_connect()
+            try_connect_done.set()
+            return result
+
+    with patch("homeassistant.components.esphome.ReconnectLogic", MockReconnectLogic):
         assert await hass.config_entries.async_setup(entry.entry_id)
         await try_connect_done.wait()
 

tests/components/esphome/test_button.py

@@ -0,0 +1,54 @@
+"""Test ESPHome buttones."""
+
+
+from unittest.mock import call
+
+from aioesphomeapi import APIClient, ButtonInfo
+
+from homeassistant.components.button import (
+    DOMAIN as BUTTON_DOMAIN,
+    SERVICE_PRESS,
+)
+from homeassistant.const import ATTR_ENTITY_ID, STATE_UNAVAILABLE, STATE_UNKNOWN
+from homeassistant.core import HomeAssistant
+
+
+async def test_button_generic_entity(
+    hass: HomeAssistant, mock_client: APIClient, mock_esphome_device
+) -> None:
+    """Test a generic button entity."""
+    entity_info = [
+        ButtonInfo(
+            object_id="mybutton",
+            key=1,
+            name="my button",
+            unique_id="my_button",
+        )
+    ]
+    states = []
+    user_service = []
+    mock_device = await mock_esphome_device(
+        mock_client=mock_client,
+        entity_info=entity_info,
+        user_service=user_service,
+        states=states,
+    )
+    state = hass.states.get("button.test_my_button")
+    assert state is not None
+    assert state.state == STATE_UNKNOWN
+
+    await hass.services.async_call(
+        BUTTON_DOMAIN,
+        SERVICE_PRESS,
+        {ATTR_ENTITY_ID: "button.test_my_button"},
+        blocking=True,
+    )
+    mock_client.button_command.assert_has_calls([call(1)])
+    state = hass.states.get("button.test_my_button")
+    assert state is not None
+    assert state.state != STATE_UNKNOWN
+
+    await mock_device.mock_disconnect(False)
+    state = hass.states.get("button.test_my_button")
+    assert state is not None
+    assert state.state == STATE_UNAVAILABLE