diff --git a/tests/components/esphome/test_entity.py b/tests/components/esphome/test_entity.py index 03fd21c32f8..609af788ceb 100644 --- a/tests/components/esphome/test_entity.py +++ b/tests/components/esphome/test_entity.py @@ -1,4 +1,5 @@ """Test ESPHome binary sensors.""" +import asyncio from collections.abc import Awaitable, Callable from typing import Any from unittest.mock import AsyncMock @@ -21,8 +22,13 @@ from homeassistant.const import ( STATE_ON, STATE_UNAVAILABLE, ) -from homeassistant.core import HomeAssistant +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import entity_registry as er +from homeassistant.helpers.event import ( + EventStateChangedData, + async_track_state_change_event, +) +from homeassistant.helpers.typing import EventType from .conftest import MockESPHomeDevice @@ -215,7 +221,19 @@ async def test_entities_removed_after_reload( ) assert await hass.config_entries.async_setup(entry.entry_id) + on_future = hass.loop.create_future() + + @callback + def _async_wait_for_on(event: EventType[EventStateChangedData]) -> None: + if event.data["new_state"].state == STATE_ON: + on_future.set_result(None) + + async_track_state_change_event( + hass, ["binary_sensor.test_mybinary_sensor"], _async_wait_for_on + ) await hass.async_block_till_done() + async with asyncio.timeout(2): + await on_future assert mock_device.entry.entry_id == entry_id state = hass.states.get("binary_sensor.test_mybinary_sensor")