Use mock_restore_cache in mqtt tests (#77297)

tests/components/mqtt/test_binary_sensor.py

@@ -53,6 +53,7 @@ from tests.common import (
     assert_setup_component,
     async_fire_mqtt_message,
     async_fire_time_changed,
+    mock_restore_cache,
 )
 
 DEFAULT_CONFIG = {
@@ -1062,10 +1063,9 @@ async def test_skip_restoring_state_with_over_due_expire_trigger(
         {},
         last_changed=datetime.fromisoformat("2022-02-02 12:01:35+01:00"),
     )
-    with patch(
+    mock_restore_cache(hass, (fake_state,))
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
+
-        return_value=fake_state,
+    with assert_setup_component(1, domain):
-    ), assert_setup_component(1, domain):
         assert await async_setup_component(hass, domain, {domain: config3})
         await hass.async_block_till_done()
         await mqtt_mock_entry_with_yaml_config()

tests/components/mqtt/test_init.py

@@ -43,6 +43,7 @@ from tests.common import (
     async_fire_time_changed,
     mock_device_registry,
     mock_registry,
+    mock_restore_cache,
 )
 from tests.testing_config.custom_components.test.sensor import DEVICE_CLASSES
 
@@ -284,12 +285,9 @@ async def test_command_template_variables(hass, mqtt_mock_entry_with_yaml_config
     """Test the rendering of entity variables."""
     topic = "test/select"
 
-    fake_state = ha.State("select.test", "milk")
+    fake_state = ha.State("select.test_select", "milk")
+    mock_restore_cache(hass, (fake_state,))
 
-    with patch(
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
-        return_value=fake_state,
-    ):
     assert await async_setup_component(
         hass,
         "select",

tests/components/mqtt/test_light.py

@@ -226,7 +226,11 @@ from .test_common import (
     help_test_update_with_json_attrs_not_dict,
 )
 
-from tests.common import assert_setup_component, async_fire_mqtt_message
+from tests.common import (
+    assert_setup_component,
+    async_fire_mqtt_message,
+    mock_restore_cache,
+)
 from tests.components.light import common
 
 DEFAULT_CONFIG = {
@@ -792,10 +796,9 @@ async def test_sending_mqtt_commands_and_optimistic(
             "color_mode": "hs",
         },
     )
-    with patch(
+    mock_restore_cache(hass, (fake_state,))
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
+
-        return_value=fake_state,
+    with assert_setup_component(1, light.DOMAIN):
-    ), assert_setup_component(1, light.DOMAIN):
         assert await async_setup_component(hass, light.DOMAIN, config)
         await hass.async_block_till_done()
         mqtt_mock = await mqtt_mock_entry_with_yaml_config()

tests/components/mqtt/test_light_json.py

@@ -124,7 +124,7 @@ from .test_common import (
     help_test_update_with_json_attrs_not_dict,
 )
 
-from tests.common import async_fire_mqtt_message
+from tests.common import async_fire_mqtt_message, mock_restore_cache
 from tests.components.light import common
 
 DEFAULT_CONFIG = {
@@ -614,11 +614,8 @@ async def test_sending_mqtt_commands_and_optimistic(
             "color_temp": 100,
         },
     )
+    mock_restore_cache(hass, (fake_state,))
 
-    with patch(
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
-        return_value=fake_state,
-    ):
     assert await async_setup_component(
         hass,
         light.DOMAIN,
@@ -754,11 +751,8 @@ async def test_sending_mqtt_commands_and_optimistic2(
             "hs_color": [100, 100],
         },
     )
+    mock_restore_cache(hass, (fake_state,))
 
-    with patch(
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
-        return_value=fake_state,
-    ):
     assert await async_setup_component(
         hass,
         light.DOMAIN,

tests/components/mqtt/test_light_template.py

@@ -74,7 +74,11 @@ from .test_common import (
     help_test_update_with_json_attrs_not_dict,
 )
 
-from tests.common import assert_setup_component, async_fire_mqtt_message
+from tests.common import (
+    assert_setup_component,
+    async_fire_mqtt_message,
+    mock_restore_cache,
+)
 from tests.components.light import common
 
 DEFAULT_CONFIG = {
@@ -334,11 +338,9 @@ async def test_sending_mqtt_commands_and_optimistic(
             "color_temp": 100,
         },
     )
+    mock_restore_cache(hass, (fake_state,))
 
-    with patch(
+    with assert_setup_component(1, light.DOMAIN):
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
-        return_value=fake_state,
-    ), assert_setup_component(1, light.DOMAIN):
         assert await async_setup_component(
             hass,
             light.DOMAIN,

tests/components/mqtt/test_scene.py

@@ -25,6 +25,8 @@ from .test_common import (
     help_test_unload_config_entry_with_platform,
 )
 
+from tests.common import mock_restore_cache
+
 DEFAULT_CONFIG = {
     scene.DOMAIN: {
         "platform": "mqtt",
@@ -45,11 +47,8 @@ def scene_platform_only():
 async def test_sending_mqtt_commands(hass, mqtt_mock_entry_with_yaml_config):
     """Test the sending MQTT commands."""
     fake_state = ha.State("scene.test", STATE_UNKNOWN)
+    mock_restore_cache(hass, (fake_state,))
 
-    with patch(
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
-        return_value=fake_state,
-    ):
     assert await async_setup_component(
         hass,
         scene.DOMAIN,

tests/components/mqtt/test_select.py

@@ -53,7 +53,7 @@ from .test_common import (
     help_test_update_with_json_attrs_not_dict,
 )
 
-from tests.common import async_fire_mqtt_message
+from tests.common import async_fire_mqtt_message, mock_restore_cache
 
 DEFAULT_CONFIG = {
     select.DOMAIN: {
@@ -152,12 +152,9 @@ async def test_run_select_service_optimistic(hass, mqtt_mock_entry_with_yaml_con
     """Test that set_value service works in optimistic mode."""
     topic = "test/select"
 
-    fake_state = ha.State("select.test", "milk")
+    fake_state = ha.State("select.test_select", "milk")
+    mock_restore_cache(hass, (fake_state,))
 
-    with patch(
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
-        return_value=fake_state,
-    ):
     assert await async_setup_component(
         hass,
         select.DOMAIN,
@@ -196,12 +193,9 @@ async def test_run_select_service_optimistic_with_command_template(
     """Test that set_value service works in optimistic mode and with a command_template."""
     topic = "test/select"
 
-    fake_state = ha.State("select.test", "milk")
+    fake_state = ha.State("select.test_select", "milk")
+    mock_restore_cache(hass, (fake_state,))
 
-    with patch(
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
-        return_value=fake_state,
-    ):
     assert await async_setup_component(
         hass,
         select.DOMAIN,

tests/components/mqtt/test_sensor.py

@@ -67,6 +67,7 @@ from tests.common import (
     assert_setup_component,
     async_fire_mqtt_message,
     async_fire_time_changed,
+    mock_restore_cache_with_extra_data,
 )
 
 DEFAULT_CONFIG = {
@@ -1160,15 +1161,9 @@ async def test_skip_restoring_state_with_over_due_expire_trigger(
         last_changed=datetime.fromisoformat("2022-02-02 12:01:35+01:00"),
     )
     fake_extra_data = MagicMock()
-    with patch(
+    mock_restore_cache_with_extra_data(hass, ((fake_state, fake_extra_data),))
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state",
+
-        return_value=fake_state,
+    with assert_setup_component(1, domain):
-    ), patch(
-        "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_extra_data",
-        return_value=fake_extra_data,
-    ), assert_setup_component(
-        1, domain
-    ):
         assert await async_setup_component(hass, domain, {domain: config3})
         await hass.async_block_till_done()
         await mqtt_mock_entry_with_yaml_config()

tests/components/mqtt/test_switch.py

@@ -47,7 +47,7 @@ from .test_common import (
     help_test_update_with_json_attrs_not_dict,
 )
 
-from tests.common import async_fire_mqtt_message
+from tests.common import async_fire_mqtt_message, mock_restore_cache
 from tests.components.switch import common
 
 DEFAULT_CONFIG = {
@@ -108,11 +108,8 @@ async def test_sending_mqtt_commands_and_optimistic(
 ):
     """Test the sending MQTT commands in optimistic mode."""
     fake_state = ha.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,
-    ):
     assert await async_setup_component(
         hass,
         switch.DOMAIN,