Create dataclass to mock entry setup in Broadlink tests (#50134)

tests/components/broadlink/__init__.py

@@ -1,4 +1,5 @@
 """Tests for the Broadlink integration."""
+from dataclasses import dataclass
 from unittest.mock import MagicMock, patch
 
 from homeassistant.components.broadlink.const import DOMAIN
@@ -70,6 +71,15 @@ BROADLINK_DEVICES = {
 }
 
 
+@dataclass
+class MockSetup:
+    """Representation of a mock setup."""
+
+    api: MagicMock
+    entry: MockConfigEntry
+    factory: MagicMock
+
+
 class BroadlinkDevice:
     """Representation of a Broadlink device."""
 
@@ -96,11 +106,11 @@ class BroadlinkDevice:
         with patch(
             "homeassistant.components.broadlink.device.blk.gendevice",
             return_value=mock_api,
-        ):
+        ) as mock_factory:
             await hass.config_entries.async_setup(mock_entry.entry_id)
             await hass.async_block_till_done()
 
-        return mock_api, mock_entry
+        return MockSetup(mock_api, mock_entry, mock_factory)
 
     def get_mock_api(self):
         """Return a mock device (API)."""

tests/components/broadlink/test_device.py

@@ -12,6 +12,8 @@ from . import get_device
 
 from tests.common import mock_device_registry, mock_registry
 
+DEVICE_FACTORY = "homeassistant.components.broadlink.device.blk.gendevice"
+
 
 async def test_device_setup(hass):
     """Test a successful setup."""
@@ -22,13 +24,15 @@ async def test_device_setup(hass):
     ) as mock_forward, patch.object(
         hass.config_entries.flow, "async_init"
     ) as mock_init:
-        mock_api, mock_entry = await device.setup_entry(hass)
+        mock_setup = await device.setup_entry(hass)
+
+    assert mock_setup.entry.state is ConfigEntryState.LOADED
+    assert mock_setup.api.auth.call_count == 1
+    assert mock_setup.api.get_fwversion.call_count == 1
+    assert mock_setup.factory.call_count == 1
 
-    assert mock_entry.state == ConfigEntryState.LOADED
-    assert mock_api.auth.call_count == 1
-    assert mock_api.get_fwversion.call_count == 1
     forward_entries = {c[1][1] for c in mock_forward.mock_calls}
-    domains = get_domains(mock_api.type)
+    domains = get_domains(mock_setup.api.type)
     assert mock_forward.call_count == len(domains)
     assert forward_entries == domains
     assert mock_init.call_count == 0
@@ -45,10 +49,10 @@ async def test_device_setup_authentication_error(hass):
     ) as mock_forward, patch.object(
         hass.config_entries.flow, "async_init"
     ) as mock_init:
-        mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+        mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
-    assert mock_entry.state == ConfigEntryState.SETUP_ERROR
+    assert mock_setup.entry.state is ConfigEntryState.SETUP_ERROR
-    assert mock_api.auth.call_count == 1
+    assert mock_setup.api.auth.call_count == 1
     assert mock_forward.call_count == 0
     assert mock_init.call_count == 1
     assert mock_init.mock_calls[0][2]["context"]["source"] == "reauth"
@@ -69,10 +73,10 @@ async def test_device_setup_network_timeout(hass):
     ) as mock_forward, patch.object(
         hass.config_entries.flow, "async_init"
     ) as mock_init:
-        mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+        mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
-    assert mock_entry.state is ConfigEntryState.SETUP_RETRY
+    assert mock_setup.entry.state is ConfigEntryState.SETUP_RETRY
-    assert mock_api.auth.call_count == 1
+    assert mock_setup.api.auth.call_count == 1
     assert mock_forward.call_count == 0
     assert mock_init.call_count == 0
 
@@ -88,10 +92,10 @@ async def test_device_setup_os_error(hass):
     ) as mock_forward, patch.object(
         hass.config_entries.flow, "async_init"
     ) as mock_init:
-        mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+        mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
-    assert mock_entry.state is ConfigEntryState.SETUP_RETRY
+    assert mock_setup.entry.state is ConfigEntryState.SETUP_RETRY
-    assert mock_api.auth.call_count == 1
+    assert mock_setup.api.auth.call_count == 1
     assert mock_forward.call_count == 0
     assert mock_init.call_count == 0
 
@@ -107,10 +111,10 @@ async def test_device_setup_broadlink_exception(hass):
     ) as mock_forward, patch.object(
         hass.config_entries.flow, "async_init"
     ) as mock_init:
-        mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+        mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
-    assert mock_entry.state is ConfigEntryState.SETUP_ERROR
+    assert mock_setup.entry.state is ConfigEntryState.SETUP_ERROR
-    assert mock_api.auth.call_count == 1
+    assert mock_setup.api.auth.call_count == 1
     assert mock_forward.call_count == 0
     assert mock_init.call_count == 0
 
@@ -126,11 +130,11 @@ async def test_device_setup_update_network_timeout(hass):
     ) as mock_forward, patch.object(
         hass.config_entries.flow, "async_init"
     ) as mock_init:
-        mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+        mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
-    assert mock_entry.state is ConfigEntryState.SETUP_RETRY
+    assert mock_setup.entry.state is ConfigEntryState.SETUP_RETRY
-    assert mock_api.auth.call_count == 1
+    assert mock_setup.api.auth.call_count == 1
-    assert mock_api.check_sensors.call_count == 1
+    assert mock_setup.api.check_sensors.call_count == 1
     assert mock_forward.call_count == 0
     assert mock_init.call_count == 0
 
@@ -149,11 +153,12 @@ async def test_device_setup_update_authorization_error(hass):
     ) as mock_forward, patch.object(
         hass.config_entries.flow, "async_init"
     ) as mock_init:
-        mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+        mock_setup = await device.setup_entry(hass, mock_api=mock_api)
+
+    assert mock_setup.entry.state is ConfigEntryState.LOADED
+    assert mock_setup.api.auth.call_count == 2
+    assert mock_setup.api.check_sensors.call_count == 2
 
-    assert mock_entry.state is ConfigEntryState.LOADED
-    assert mock_api.auth.call_count == 2
-    assert mock_api.check_sensors.call_count == 2
     forward_entries = {c[1][1] for c in mock_forward.mock_calls}
     domains = get_domains(mock_api.type)
     assert mock_forward.call_count == len(domains)
@@ -173,11 +178,11 @@ async def test_device_setup_update_authentication_error(hass):
     ) as mock_forward, patch.object(
         hass.config_entries.flow, "async_init"
     ) as mock_init:
-        mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+        mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
-    assert mock_entry.state is ConfigEntryState.SETUP_RETRY
+    assert mock_setup.entry.state is ConfigEntryState.SETUP_RETRY
-    assert mock_api.auth.call_count == 2
+    assert mock_setup.api.auth.call_count == 2
-    assert mock_api.check_sensors.call_count == 1
+    assert mock_setup.api.check_sensors.call_count == 1
     assert mock_forward.call_count == 0
     assert mock_init.call_count == 1
     assert mock_init.mock_calls[0][2]["context"]["source"] == "reauth"
@@ -198,11 +203,11 @@ async def test_device_setup_update_broadlink_exception(hass):
     ) as mock_forward, patch.object(
         hass.config_entries.flow, "async_init"
     ) as mock_init:
-        mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+        mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
-    assert mock_entry.state is ConfigEntryState.SETUP_RETRY
+    assert mock_setup.entry.state is ConfigEntryState.SETUP_RETRY
-    assert mock_api.auth.call_count == 1
+    assert mock_setup.api.auth.call_count == 1
-    assert mock_api.check_sensors.call_count == 1
+    assert mock_setup.api.check_sensors.call_count == 1
     assert mock_forward.call_count == 0
     assert mock_init.call_count == 0
 
@@ -214,11 +219,11 @@ async def test_device_setup_get_fwversion_broadlink_exception(hass):
     mock_api.get_fwversion.side_effect = blke.BroadlinkException()
 
     with patch.object(hass.config_entries, "async_forward_entry_setup") as mock_forward:
-        mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+        mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
-    assert mock_entry.state is ConfigEntryState.LOADED
+    assert mock_setup.entry.state is ConfigEntryState.LOADED
     forward_entries = {c[1][1] for c in mock_forward.mock_calls}
-    domains = get_domains(mock_api.type)
+    domains = get_domains(mock_setup.api.type)
     assert mock_forward.call_count == len(domains)
     assert forward_entries == domains
 
@@ -230,11 +235,11 @@ async def test_device_setup_get_fwversion_os_error(hass):
     mock_api.get_fwversion.side_effect = OSError()
 
     with patch.object(hass.config_entries, "async_forward_entry_setup") as mock_forward:
-        _, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+        mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
-    assert mock_entry.state is ConfigEntryState.LOADED
+    assert mock_setup.entry.state is ConfigEntryState.LOADED
     forward_entries = {c[1][1] for c in mock_forward.mock_calls}
-    domains = get_domains(mock_api.type)
+    domains = get_domains(mock_setup.api.type)
     assert mock_forward.call_count == len(domains)
     assert forward_entries == domains
 
@@ -246,12 +251,14 @@ async def test_device_setup_registry(hass):
     device_registry = mock_device_registry(hass)
     entity_registry = mock_registry(hass)
 
-    _, mock_entry = await device.setup_entry(hass)
+    mock_setup = await device.setup_entry(hass)
     await hass.async_block_till_done()
 
     assert len(device_registry.devices) == 1
 
-    device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
+    device_entry = device_registry.async_get_device(
+        {(DOMAIN, mock_setup.entry.unique_id)}
+    )
     assert device_entry.identifiers == {(DOMAIN, device.mac)}
     assert device_entry.name == device.name
     assert device_entry.model == device.model
@@ -267,16 +274,16 @@ async def test_device_unload_works(hass):
     device = get_device("Office")
 
     with patch.object(hass.config_entries, "async_forward_entry_setup"):
-        mock_api, mock_entry = await device.setup_entry(hass)
+        mock_setup = await device.setup_entry(hass)
 
     with patch.object(
         hass.config_entries, "async_forward_entry_unload", return_value=True
     ) as mock_forward:
-        await hass.config_entries.async_unload(mock_entry.entry_id)
+        await hass.config_entries.async_unload(mock_setup.entry.entry_id)
 
-    assert mock_entry.state is ConfigEntryState.NOT_LOADED
+    assert mock_setup.entry.state is ConfigEntryState.NOT_LOADED
     forward_entries = {c[1][1] for c in mock_forward.mock_calls}
-    domains = get_domains(mock_api.type)
+    domains = get_domains(mock_setup.api.type)
     assert mock_forward.call_count == len(domains)
     assert forward_entries == domains
 
@@ -290,14 +297,14 @@ async def test_device_unload_authentication_error(hass):
     with patch.object(hass.config_entries, "async_forward_entry_setup"), patch.object(
         hass.config_entries.flow, "async_init"
     ):
-        _, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+        mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
     with patch.object(
         hass.config_entries, "async_forward_entry_unload", return_value=True
     ) as mock_forward:
-        await hass.config_entries.async_unload(mock_entry.entry_id)
+        await hass.config_entries.async_unload(mock_setup.entry.entry_id)
 
-    assert mock_entry.state is ConfigEntryState.NOT_LOADED
+    assert mock_setup.entry.state is ConfigEntryState.NOT_LOADED
     assert mock_forward.call_count == 0
 
 
@@ -308,14 +315,14 @@ async def test_device_unload_update_failed(hass):
     mock_api.check_sensors.side_effect = blke.NetworkTimeoutError()
 
     with patch.object(hass.config_entries, "async_forward_entry_setup"):
-        _, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+        mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
     with patch.object(
         hass.config_entries, "async_forward_entry_unload", return_value=True
     ) as mock_forward:
-        await hass.config_entries.async_unload(mock_entry.entry_id)
+        await hass.config_entries.async_unload(mock_setup.entry.entry_id)
 
-    assert mock_entry.state is ConfigEntryState.NOT_LOADED
+    assert mock_setup.entry.state is ConfigEntryState.NOT_LOADED
     assert mock_forward.call_count == 0
 
 
@@ -326,16 +333,16 @@ async def test_device_update_listener(hass):
     device_registry = mock_device_registry(hass)
     entity_registry = mock_registry(hass)
 
-    mock_api, mock_entry = await device.setup_entry(hass)
+    mock_setup = await device.setup_entry(hass)
     await hass.async_block_till_done()
 
-    with patch(
+    with patch(DEVICE_FACTORY, return_value=mock_setup.api):
-        "homeassistant.components.broadlink.device.blk.gendevice", return_value=mock_api
+        hass.config_entries.async_update_entry(mock_setup.entry, title="New Name")
-    ):
-        hass.config_entries.async_update_entry(mock_entry, title="New Name")
         await hass.async_block_till_done()
 
-    device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
+    device_entry = device_registry.async_get_device(
+        {(DOMAIN, mock_setup.entry.unique_id)}
+    )
     assert device_entry.name == "New Name"
     for entry in async_entries_for_device(entity_registry, device_entry.id):
         assert entry.original_name.startswith("New Name")

tests/components/broadlink/test_heartbeat.py

@@ -72,10 +72,10 @@ async def test_heartbeat_unload(hass):
     """Test that the heartbeat is deactivated when the last config entry is removed."""
     device = get_device("Office")
 
-    _, mock_entry = await device.setup_entry(hass)
+    mock_setup = await device.setup_entry(hass)
     await hass.async_block_till_done()
 
-    await hass.config_entries.async_remove(mock_entry.entry_id)
+    await hass.config_entries.async_remove(mock_setup.entry.entry_id)
     await hass.async_block_till_done()
 
     with patch(DEVICE_PING) as mock_ping:
@@ -91,11 +91,11 @@ async def test_heartbeat_do_not_unload(hass):
     device_a = get_device("Office")
     device_b = get_device("Bedroom")
 
-    _, mock_entry_a = await device_a.setup_entry(hass)
+    mock_setup = await device_a.setup_entry(hass)
     await device_b.setup_entry(hass)
     await hass.async_block_till_done()
 
-    await hass.config_entries.async_remove(mock_entry_a.entry_id)
+    await hass.config_entries.async_remove(mock_setup.entry.entry_id)
     await hass.async_block_till_done()
 
     with patch(DEVICE_PING) as mock_ping:

tests/components/broadlink/test_remote.py

@@ -28,10 +28,10 @@ async def test_remote_setup_works(hass):
     for device in map(get_device, REMOTE_DEVICES):
         device_registry = mock_device_registry(hass)
         entity_registry = mock_registry(hass)
-        mock_api, mock_entry = await device.setup_entry(hass)
+        mock_setup = await device.setup_entry(hass)
 
         device_entry = device_registry.async_get_device(
-            {(DOMAIN, mock_entry.unique_id)}
+            {(DOMAIN, mock_setup.entry.unique_id)}
         )
         entries = async_entries_for_device(entity_registry, device_entry.id)
         remotes = {entry for entry in entries if entry.domain == REMOTE_DOMAIN}
@@ -40,7 +40,7 @@ async def test_remote_setup_works(hass):
         remote = remotes.pop()
         assert remote.original_name == f"{device.name} Remote"
         assert hass.states.get(remote.entity_id).state == STATE_ON
-        assert mock_api.auth.call_count == 1
+        assert mock_setup.api.auth.call_count == 1
 
 
 async def test_remote_send_command(hass):
@@ -48,10 +48,10 @@ async def test_remote_send_command(hass):
     for device in map(get_device, REMOTE_DEVICES):
         device_registry = mock_device_registry(hass)
         entity_registry = mock_registry(hass)
-        mock_api, mock_entry = await device.setup_entry(hass)
+        mock_setup = await device.setup_entry(hass)
 
         device_entry = device_registry.async_get_device(
-            {(DOMAIN, mock_entry.unique_id)}
+            {(DOMAIN, mock_setup.entry.unique_id)}
         )
         entries = async_entries_for_device(entity_registry, device_entry.id)
         remotes = {entry for entry in entries if entry.domain == REMOTE_DOMAIN}
@@ -65,9 +65,9 @@ async def test_remote_send_command(hass):
             blocking=True,
         )
 
-        assert mock_api.send_data.call_count == 1
+        assert mock_setup.api.send_data.call_count == 1
-        assert mock_api.send_data.call_args == call(b64decode(IR_PACKET))
+        assert mock_setup.api.send_data.call_args == call(b64decode(IR_PACKET))
-        assert mock_api.auth.call_count == 1
+        assert mock_setup.api.auth.call_count == 1
 
 
 async def test_remote_turn_off_turn_on(hass):
@@ -75,10 +75,10 @@ async def test_remote_turn_off_turn_on(hass):
     for device in map(get_device, REMOTE_DEVICES):
         device_registry = mock_device_registry(hass)
         entity_registry = mock_registry(hass)
-        mock_api, mock_entry = await device.setup_entry(hass)
+        mock_setup = await device.setup_entry(hass)
 
         device_entry = device_registry.async_get_device(
-            {(DOMAIN, mock_entry.unique_id)}
+            {(DOMAIN, mock_setup.entry.unique_id)}
         )
         entries = async_entries_for_device(entity_registry, device_entry.id)
         remotes = {entry for entry in entries if entry.domain == REMOTE_DOMAIN}
@@ -99,7 +99,7 @@ async def test_remote_turn_off_turn_on(hass):
             {"entity_id": remote.entity_id, "command": "b64:" + IR_PACKET},
             blocking=True,
         )
-        assert mock_api.send_data.call_count == 0
+        assert mock_setup.api.send_data.call_count == 0
 
         await hass.services.async_call(
             REMOTE_DOMAIN,
@@ -115,6 +115,6 @@ async def test_remote_turn_off_turn_on(hass):
             {"entity_id": remote.entity_id, "command": "b64:" + IR_PACKET},
             blocking=True,
         )
-        assert mock_api.send_data.call_count == 1
+        assert mock_setup.api.send_data.call_count == 1
-        assert mock_api.send_data.call_args == call(b64decode(IR_PACKET))
+        assert mock_setup.api.send_data.call_args == call(b64decode(IR_PACKET))
-        assert mock_api.auth.call_count == 1
+        assert mock_setup.api.auth.call_count == 1

tests/components/broadlink/test_sensors.py

@@ -22,10 +22,12 @@ async def test_a1_sensor_setup(hass):
     device_registry = mock_device_registry(hass)
     entity_registry = mock_registry(hass)
 
-    mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+    mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
     assert mock_api.check_sensors_raw.call_count == 1
-    device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
+    device_entry = device_registry.async_get_device(
+        {(DOMAIN, mock_setup.entry.unique_id)}
+    )
     entries = async_entries_for_device(entity_registry, device_entry.id)
     sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 5
@@ -58,14 +60,16 @@ async def test_a1_sensor_update(hass):
     device_registry = mock_device_registry(hass)
     entity_registry = mock_registry(hass)
 
-    mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+    mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
-    device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
+    device_entry = device_registry.async_get_device(
+        {(DOMAIN, mock_setup.entry.unique_id)}
+    )
     entries = async_entries_for_device(entity_registry, device_entry.id)
     sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 5
 
-    mock_api.check_sensors_raw.return_value = {
+    mock_setup.api.check_sensors_raw.return_value = {
         "temperature": 22.5,
         "humidity": 47.4,
         "air_quality": 2,
@@ -75,7 +79,7 @@ async def test_a1_sensor_update(hass):
     await hass.helpers.entity_component.async_update_entity(
         next(iter(sensors)).entity_id
     )
-    assert mock_api.check_sensors_raw.call_count == 2
+    assert mock_setup.api.check_sensors_raw.call_count == 2
 
     sensors_and_states = {
         (sensor.original_name, hass.states.get(sensor.entity_id).state)
@@ -99,10 +103,12 @@ async def test_rm_pro_sensor_setup(hass):
     device_registry = mock_device_registry(hass)
     entity_registry = mock_registry(hass)
 
-    mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+    mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
     assert mock_api.check_sensors.call_count == 1
-    device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
+    device_entry = device_registry.async_get_device(
+        {(DOMAIN, mock_setup.entry.unique_id)}
+    )
     entries = async_entries_for_device(entity_registry, device_entry.id)
     sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 1
@@ -123,18 +129,20 @@ async def test_rm_pro_sensor_update(hass):
     device_registry = mock_device_registry(hass)
     entity_registry = mock_registry(hass)
 
-    mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+    mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
-    device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
+    device_entry = device_registry.async_get_device(
+        {(DOMAIN, mock_setup.entry.unique_id)}
+    )
     entries = async_entries_for_device(entity_registry, device_entry.id)
     sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 1
 
-    mock_api.check_sensors.return_value = {"temperature": 25.8}
+    mock_setup.api.check_sensors.return_value = {"temperature": 25.8}
     await hass.helpers.entity_component.async_update_entity(
         next(iter(sensors)).entity_id
     )
-    assert mock_api.check_sensors.call_count == 2
+    assert mock_setup.api.check_sensors.call_count == 2
 
     sensors_and_states = {
         (sensor.original_name, hass.states.get(sensor.entity_id).state)
@@ -155,18 +163,20 @@ async def test_rm_pro_filter_crazy_temperature(hass):
     device_registry = mock_device_registry(hass)
     entity_registry = mock_registry(hass)
 
-    mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+    mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
-    device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
+    device_entry = device_registry.async_get_device(
+        {(DOMAIN, mock_setup.entry.unique_id)}
+    )
     entries = async_entries_for_device(entity_registry, device_entry.id)
     sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 1
 
-    mock_api.check_sensors.return_value = {"temperature": -7}
+    mock_setup.api.check_sensors.return_value = {"temperature": -7}
     await hass.helpers.entity_component.async_update_entity(
         next(iter(sensors)).entity_id
     )
-    assert mock_api.check_sensors.call_count == 2
+    assert mock_setup.api.check_sensors.call_count == 2
 
     sensors_and_states = {
         (sensor.original_name, hass.states.get(sensor.entity_id).state)
@@ -184,10 +194,12 @@ async def test_rm_mini3_no_sensor(hass):
     device_registry = mock_device_registry(hass)
     entity_registry = mock_registry(hass)
 
-    mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+    mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
     assert mock_api.check_sensors.call_count <= 1
-    device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
+    device_entry = device_registry.async_get_device(
+        {(DOMAIN, mock_setup.entry.unique_id)}
+    )
     entries = async_entries_for_device(entity_registry, device_entry.id)
     sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 0
@@ -202,10 +214,12 @@ async def test_rm4_pro_hts2_sensor_setup(hass):
     device_registry = mock_device_registry(hass)
     entity_registry = mock_registry(hass)
 
-    mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+    mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
     assert mock_api.check_sensors.call_count == 1
-    device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
+    device_entry = device_registry.async_get_device(
+        {(DOMAIN, mock_setup.entry.unique_id)}
+    )
     entries = async_entries_for_device(entity_registry, device_entry.id)
     sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 2
@@ -229,18 +243,20 @@ async def test_rm4_pro_hts2_sensor_update(hass):
     device_registry = mock_device_registry(hass)
     entity_registry = mock_registry(hass)
 
-    mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+    mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
-    device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
+    device_entry = device_registry.async_get_device(
+        {(DOMAIN, mock_setup.entry.unique_id)}
+    )
     entries = async_entries_for_device(entity_registry, device_entry.id)
     sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 2
 
-    mock_api.check_sensors.return_value = {"temperature": 16.8, "humidity": 34.0}
+    mock_setup.api.check_sensors.return_value = {"temperature": 16.8, "humidity": 34.0}
     await hass.helpers.entity_component.async_update_entity(
         next(iter(sensors)).entity_id
     )
-    assert mock_api.check_sensors.call_count == 2
+    assert mock_setup.api.check_sensors.call_count == 2
 
     sensors_and_states = {
         (sensor.original_name, hass.states.get(sensor.entity_id).state)
@@ -261,10 +277,12 @@ async def test_rm4_pro_no_sensor(hass):
     device_registry = mock_device_registry(hass)
     entity_registry = mock_registry(hass)
 
-    mock_api, mock_entry = await device.setup_entry(hass, mock_api=mock_api)
+    mock_setup = await device.setup_entry(hass, mock_api=mock_api)
 
     assert mock_api.check_sensors.call_count <= 1
-    device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
+    device_entry = device_registry.async_get_device(
+        {(DOMAIN, mock_setup.entry.unique_id)}
+    )
     entries = async_entries_for_device(entity_registry, device_entry.id)
     sensors = {entry for entry in entries if entry.domain == SENSOR_DOMAIN}
     assert len(sensors) == 0