From 5bf80a0f6d0a2630851aaf29bb01ad9a5399e1e9 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Wed, 16 Aug 2023 11:05:22 -0500
Subject: [PATCH] Make ESPHome deep sleep tests more robust (#98535)

---
 tests/components/esphome/test_entity.py | 50 ++++++++++++++++++++++++-
 1 file changed, 48 insertions(+), 2 deletions(-)

diff --git a/tests/components/esphome/test_entity.py b/tests/components/esphome/test_entity.py
index ac121a93eff..fdc57b2dc24 100644
--- a/tests/components/esphome/test_entity.py
+++ b/tests/components/esphome/test_entity.py
@@ -8,10 +8,12 @@ from aioesphomeapi import (
     BinarySensorState,
     EntityInfo,
     EntityState,
+    SensorInfo,
+    SensorState,
     UserService,
 )
 
-from homeassistant.const import ATTR_RESTORED, STATE_ON, STATE_UNAVAILABLE
+from homeassistant.const import ATTR_RESTORED, STATE_OFF, STATE_ON, STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 
 from .conftest import MockESPHomeDevice
@@ -149,10 +151,17 @@ async def test_deep_sleep_device(
             name="my binary_sensor",
             unique_id="my_binary_sensor",
         ),
+        SensorInfo(
+            object_id="my_sensor",
+            key=3,
+            name="my sensor",
+            unique_id="my_sensor",
+        ),
     ]
     states = [
         BinarySensorState(key=1, state=True, missing_state=False),
         BinarySensorState(key=2, state=True, missing_state=False),
+        SensorState(key=3, state=123.0, missing_state=False),
     ]
     user_service = []
     mock_device = await mock_esphome_device(
@@ -165,12 +174,18 @@ async def test_deep_sleep_device(
     state = hass.states.get("binary_sensor.test_mybinary_sensor")
     assert state is not None
     assert state.state == STATE_ON
+    state = hass.states.get("sensor.test_my_sensor")
+    assert state is not None
+    assert state.state == "123"
 
     await mock_device.mock_disconnect(False)
     await hass.async_block_till_done()
     state = hass.states.get("binary_sensor.test_mybinary_sensor")
     assert state is not None
     assert state.state == STATE_UNAVAILABLE
+    state = hass.states.get("sensor.test_my_sensor")
+    assert state is not None
+    assert state.state == STATE_UNAVAILABLE
 
     await mock_device.mock_connect()
     await hass.async_block_till_done()
@@ -178,12 +193,43 @@ async def test_deep_sleep_device(
     state = hass.states.get("binary_sensor.test_mybinary_sensor")
     assert state is not None
     assert state.state == STATE_ON
+    state = hass.states.get("sensor.test_my_sensor")
+    assert state is not None
+    assert state.state == "123"
+
+    await mock_device.mock_disconnect(True)
+    await hass.async_block_till_done()
+    await mock_device.mock_connect()
+    await hass.async_block_till_done()
+    mock_device.set_state(BinarySensorState(key=1, state=False, missing_state=False))
+    mock_device.set_state(SensorState(key=3, state=56, missing_state=False))
+    await hass.async_block_till_done()
+    state = hass.states.get("binary_sensor.test_mybinary_sensor")
+    assert state is not None
+    assert state.state == STATE_OFF
+    state = hass.states.get("sensor.test_my_sensor")
+    assert state is not None
+    assert state.state == "56"
 
     await mock_device.mock_disconnect(True)
     await hass.async_block_till_done()
     state = hass.states.get("binary_sensor.test_mybinary_sensor")
     assert state is not None
-    assert state.state == STATE_ON
+    assert state.state == STATE_OFF
+    state = hass.states.get("sensor.test_my_sensor")
+    assert state is not None
+    assert state.state == "56"
+
+    await mock_device.mock_connect()
+    await hass.async_block_till_done()
+    await mock_device.mock_disconnect(False)
+    await hass.async_block_till_done()
+    state = hass.states.get("binary_sensor.test_mybinary_sensor")
+    assert state is not None
+    assert state.state == STATE_UNAVAILABLE
+    state = hass.states.get("sensor.test_my_sensor")
+    assert state is not None
+    assert state.state == STATE_UNAVAILABLE
 
 
 async def test_esphome_device_without_friendly_name(