From ca558545a1ce19a8f643a82a24b77c2c5cb43860 Mon Sep 17 00:00:00 2001
From: jan iversen <jan@casacondor.com>
Date: Sat, 15 May 2021 21:39:41 +0200
Subject: [PATCH] Use mock_restore_state in testing of modbus sensor (#50455)

---
 tests/components/modbus/test_modbus_sensor.py | 36 +++++++++----------
 1 file changed, 17 insertions(+), 19 deletions(-)

diff --git a/tests/components/modbus/test_modbus_sensor.py b/tests/components/modbus/test_modbus_sensor.py
index c83eab0b3d6..cb784ac46b3 100644
--- a/tests/components/modbus/test_modbus_sensor.py
+++ b/tests/components/modbus/test_modbus_sensor.py
@@ -1,6 +1,5 @@
 """The tests for the Modbus sensor component."""
 import logging
-from unittest import mock
 
 import pytest
 
@@ -42,6 +41,8 @@ from homeassistant.core import State
 
 from .conftest import ReadResult, base_config_test, base_test, prepare_service_update
 
+from tests.common import mock_restore_cache
+
 
 @pytest.mark.parametrize(
     "do_discovery, do_config",
@@ -573,24 +574,21 @@ async def test_restore_state_sensor(hass):
     sensor_name = "test_sensor"
     test_value = "117"
     config_sensor = {CONF_NAME: sensor_name, CONF_ADDRESS: 17}
-    with mock.patch(
-        "homeassistant.components.modbus.sensor.ModbusRegisterSensor.async_get_last_state"
-    ) as mock_get_last_state:
-        mock_get_last_state.return_value = State(
-            f"{SENSOR_DOMAIN}.{sensor_name}", f"{test_value}"
-        )
-
-        await base_config_test(
-            hass,
-            config_sensor,
-            sensor_name,
-            SENSOR_DOMAIN,
-            CONF_SENSORS,
-            None,
-            method_discovery=True,
-        )
-        entity_id = f"{SENSOR_DOMAIN}.{sensor_name}"
-        assert hass.states.get(entity_id).state == test_value
+    mock_restore_cache(
+        hass,
+        (State(f"{SENSOR_DOMAIN}.{sensor_name}", test_value),),
+    )
+    await base_config_test(
+        hass,
+        config_sensor,
+        sensor_name,
+        SENSOR_DOMAIN,
+        CONF_SENSORS,
+        None,
+        method_discovery=True,
+    )
+    entity_id = f"{SENSOR_DOMAIN}.{sensor_name}"
+    assert hass.states.get(entity_id).state == test_value
 
 
 @pytest.mark.parametrize(