Use mock_restore_cache in tests (#77298)

tests/components/knx/test_binary_sensor.py

@@ -1,6 +1,5 @@
 """Test KNX binary sensor."""
 from datetime import timedelta
-from unittest.mock import patch
 
 from homeassistant.components.knx.const import CONF_STATE_ADDRESS, CONF_SYNC_STATE
 from homeassistant.components.knx.schema import BinarySensorSchema
@@ -12,7 +11,11 @@ from homeassistant.util import dt
 
 from .conftest import KNXTestKit
 
-from tests.common import async_capture_events, async_fire_time_changed
+from tests.common import (
+    async_capture_events,
+    async_fire_time_changed,
+    mock_restore_cache,
+)
 
 
 async def test_binary_sensor_entity_category(hass: HomeAssistant, knx: KNXTestKit):
@@ -245,22 +248,19 @@ async def test_binary_sensor_restore_and_respond(hass, knx):
     """Test restoring KNX binary sensor state and respond to read."""
     _ADDRESS = "2/2/2"
     fake_state = State("binary_sensor.test", STATE_ON)
+    mock_restore_cache(hass, (fake_state,))
 
-    with patch(
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
-        return_value=fake_state,
-    ):
-        await knx.setup_integration(
-            {
-                BinarySensorSchema.PLATFORM: [
-                    {
-                        CONF_NAME: "test",
-                        CONF_STATE_ADDRESS: _ADDRESS,
-                        CONF_SYNC_STATE: False,
-                    },
-                ]
-            }
-        )
+    await knx.setup_integration(
+        {
+            BinarySensorSchema.PLATFORM: [
+                {
+                    CONF_NAME: "test",
+                    CONF_STATE_ADDRESS: _ADDRESS,
+                    CONF_SYNC_STATE: False,
+                },
+            ]
+        }
+    )
 
     # restored state - doesn't send telegram
     state = hass.states.get("binary_sensor.test")
@@ -277,23 +277,20 @@ async def test_binary_sensor_restore_invert(hass, knx):
     """Test restoring KNX binary sensor state with invert."""
     _ADDRESS = "2/2/2"
     fake_state = State("binary_sensor.test", STATE_ON)
+    mock_restore_cache(hass, (fake_state,))
 
-    with patch(
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
-        return_value=fake_state,
-    ):
-        await knx.setup_integration(
-            {
-                BinarySensorSchema.PLATFORM: [
-                    {
-                        CONF_NAME: "test",
-                        CONF_STATE_ADDRESS: _ADDRESS,
-                        BinarySensorSchema.CONF_INVERT: True,
-                        CONF_SYNC_STATE: False,
-                    },
-                ]
-            }
-        )
+    await knx.setup_integration(
+        {
+            BinarySensorSchema.PLATFORM: [
+                {
+                    CONF_NAME: "test",
+                    CONF_STATE_ADDRESS: _ADDRESS,
+                    BinarySensorSchema.CONF_INVERT: True,
+                    CONF_SYNC_STATE: False,
+                },
+            ]
+        }
+    )
 
     # restored state - doesn't send telegram
     state = hass.states.get("binary_sensor.test")

tests/components/knx/test_select.py

@@ -1,6 +1,4 @@
 """Test KNX select."""
-from unittest.mock import patch
-
 import pytest
 
 from homeassistant.components.knx.const import (
@@ -17,6 +15,8 @@ from homeassistant.core import HomeAssistant, State
 
 from .conftest import KNXTestKit
 
+from tests.common import mock_restore_cache
+
 
 async def test_select_dpt_2_simple(hass: HomeAssistant, knx: KNXTestKit):
     """Test simple KNX select."""
@@ -98,22 +98,19 @@ async def test_select_dpt_2_restore(hass: HomeAssistant, knx: KNXTestKit):
     test_address = "1/1/1"
     test_passive_address = "3/3/3"
     fake_state = State("select.test", "Control - On")
+    mock_restore_cache(hass, (fake_state,))
 
-    with patch(
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
-        return_value=fake_state,
-    ):
-        await knx.setup_integration(
-            {
-                SelectSchema.PLATFORM: {
-                    CONF_NAME: "test",
-                    KNX_ADDRESS: [test_address, test_passive_address],
-                    CONF_RESPOND_TO_READ: True,
-                    CONF_PAYLOAD_LENGTH: 0,
-                    SelectSchema.CONF_OPTIONS: _options,
-                }
+    await knx.setup_integration(
+        {
+            SelectSchema.PLATFORM: {
+                CONF_NAME: "test",
+                KNX_ADDRESS: [test_address, test_passive_address],
+                CONF_RESPOND_TO_READ: True,
+                CONF_PAYLOAD_LENGTH: 0,
+                SelectSchema.CONF_OPTIONS: _options,
             }
-        )
+        }
+    )
     # restored state - doesn't send telegram
     state = hass.states.get("select.test")
     assert state.state == "Control - On"

tests/components/knx/test_switch.py

@@ -1,6 +1,4 @@
 """Test KNX switch."""
-from unittest.mock import patch
-
 from homeassistant.components.knx.const import (
     CONF_RESPOND_TO_READ,
     CONF_STATE_ADDRESS,
@@ -12,6 +10,8 @@ from homeassistant.core import HomeAssistant, State
 
 from .conftest import KNXTestKit
 
+from tests.common import mock_restore_cache
+
 
 async def test_switch_simple(hass: HomeAssistant, knx: KNXTestKit):
     """Test simple KNX switch."""
@@ -115,20 +115,17 @@ async def test_switch_restore_and_respond(hass, knx):
     """Test restoring KNX switch state and respond to read."""
     _ADDRESS = "1/1/1"
     fake_state = State("switch.test", "on")
+    mock_restore_cache(hass, (fake_state,))
 
-    with patch(
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
-        return_value=fake_state,
-    ):
-        await knx.setup_integration(
-            {
-                SwitchSchema.PLATFORM: {
-                    CONF_NAME: "test",
-                    KNX_ADDRESS: _ADDRESS,
-                    CONF_RESPOND_TO_READ: True,
-                },
-            }
-        )
+    await knx.setup_integration(
+        {
+            SwitchSchema.PLATFORM: {
+                CONF_NAME: "test",
+                KNX_ADDRESS: _ADDRESS,
+                CONF_RESPOND_TO_READ: True,
+            },
+        }
+    )
 
     # restored state - doesn't send telegram
     state = hass.states.get("switch.test")

tests/components/unifi/test_switch.py

@@ -1,6 +1,5 @@
 """UniFi Network switch platform tests."""
 from copy import deepcopy
-from unittest.mock import patch
 
 from aiounifi.controller import MESSAGE_CLIENT_REMOVED, MESSAGE_DEVICE, MESSAGE_EVENT
 
@@ -32,6 +31,8 @@ from .test_controller import (
     setup_unifi_integration,
 )
 
+from tests.common import mock_restore_cache
+
 CLIENT_1 = {
     "hostname": "client_1",
     "ip": "10.0.0.1",
@@ -1249,6 +1250,7 @@ async def test_restore_client_succeed(hass, aioclient_mock):
             "poe_mode": "auto",
         },
     )
+    mock_restore_cache(hass, (fake_state,))
 
     config_entry = config_entries.ConfigEntry(
         version=1,
@@ -1269,21 +1271,17 @@ async def test_restore_client_succeed(hass, aioclient_mock):
         config_entry=config_entry,
     )
 
-    with patch(
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
-        return_value=fake_state,
-    ):
-        await setup_unifi_integration(
-            hass,
-            aioclient_mock,
-            options={
-                CONF_TRACK_CLIENTS: False,
-                CONF_TRACK_DEVICES: False,
-            },
-            clients_response=[],
-            devices_response=[POE_DEVICE],
-            clients_all_response=[POE_CLIENT],
-        )
+    await setup_unifi_integration(
+        hass,
+        aioclient_mock,
+        options={
+            CONF_TRACK_CLIENTS: False,
+            CONF_TRACK_DEVICES: False,
+        },
+        clients_response=[],
+        devices_response=[POE_DEVICE],
+        clients_all_response=[POE_CLIENT],
+    )
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1