Use an assumed switch state until refresh callback is complete (#59805)

tests/components/litterrobot/test_switch.py

@@ -1,5 +1,6 @@
 """Test the Litter-Robot switch entity."""
 from datetime import timedelta
+from unittest.mock import MagicMock
 
 import pytest
 
@@ -9,7 +10,13 @@ from homeassistant.components.switch import (
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
 )
-from homeassistant.const import ATTR_ENTITY_ID, ENTITY_CATEGORY_CONFIG, STATE_ON
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    ENTITY_CATEGORY_CONFIG,
+    STATE_OFF,
+    STATE_ON,
+)
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry
 from homeassistant.util.dt import utcnow
 
@@ -21,13 +28,13 @@ NIGHT_LIGHT_MODE_ENTITY_ID = "switch.test_night_light_mode"
 PANEL_LOCKOUT_ENTITY_ID = "switch.test_panel_lockout"
 
 
-async def test_switch(hass, mock_account):
+async def test_switch(hass: HomeAssistant, mock_account: MagicMock):
     """Tests the switch entity was set up."""
     await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
 
-    switch = hass.states.get(NIGHT_LIGHT_MODE_ENTITY_ID)
-    assert switch
-    assert switch.state == STATE_ON
+    state = hass.states.get(NIGHT_LIGHT_MODE_ENTITY_ID)
+    assert state
+    assert state.state == STATE_ON
 
     ent_reg = entity_registry.async_get(hass)
     entity_entry = ent_reg.async_get(NIGHT_LIGHT_MODE_ENTITY_ID)
@@ -42,12 +49,14 @@ async def test_switch(hass, mock_account):
         (PANEL_LOCKOUT_ENTITY_ID, "set_panel_lockout"),
     ],
 )
-async def test_on_off_commands(hass, mock_account, entity_id, robot_command):
+async def test_on_off_commands(
+    hass: HomeAssistant, mock_account: MagicMock, entity_id: str, robot_command: str
+):
     """Test sending commands to the switch."""
     await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
 
-    switch = hass.states.get(entity_id)
-    assert switch
+    state = hass.states.get(entity_id)
+    assert state
 
     data = {ATTR_ENTITY_ID: entity_id}
 
@@ -65,3 +74,6 @@ async def test_on_off_commands(hass, mock_account, entity_id, robot_command):
         future = utcnow() + timedelta(seconds=REFRESH_WAIT_TIME_SECONDS)
         async_fire_time_changed(hass, future)
         assert getattr(mock_account.robots[0], robot_command).call_count == count
+        state = hass.states.get(entity_id)
+        assert state
+        assert state.state == STATE_ON if service == SERVICE_TURN_ON else STATE_OFF