From fca79237c3dfc3e7adefff05afccca555a85d745 Mon Sep 17 00:00:00 2001 From: Felipe Martins Diel <41558831+felipediel@users.noreply.github.com> Date: Mon, 14 Jun 2021 08:28:57 -0300 Subject: [PATCH] Create dataclass to mock entry setup in Broadlink tests (#50134) --- tests/components/broadlink/__init__.py | 14 ++- tests/components/broadlink/test_device.py | 121 ++++++++++--------- tests/components/broadlink/test_heartbeat.py | 8 +- tests/components/broadlink/test_remote.py | 28 ++--- tests/components/broadlink/test_sensors.py | 70 +++++++---- 5 files changed, 138 insertions(+), 103 deletions(-) diff --git a/tests/components/broadlink/__init__.py b/tests/components/broadlink/__init__.py index 780887551f2..c65870add96 100644 --- a/tests/components/broadlink/__init__.py +++ b/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).""" diff --git a/tests/components/broadlink/test_device.py b/tests/components/broadlink/test_device.py index 2ee8f7a5218..5430af9e311 100644 --- a/tests/components/broadlink/test_device.py +++ b/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_api.auth.call_count == 1 + assert mock_setup.entry.state is ConfigEntryState.SETUP_ERROR + 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_api.auth.call_count == 1 + assert mock_setup.entry.state is ConfigEntryState.SETUP_RETRY + 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_api.auth.call_count == 1 + assert mock_setup.entry.state is ConfigEntryState.SETUP_RETRY + 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_api.auth.call_count == 1 + assert mock_setup.entry.state is ConfigEntryState.SETUP_ERROR + 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_api.auth.call_count == 1 - assert mock_api.check_sensors.call_count == 1 + assert mock_setup.entry.state is ConfigEntryState.SETUP_RETRY + assert mock_setup.api.auth.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_api.auth.call_count == 2 - assert mock_api.check_sensors.call_count == 1 + assert mock_setup.entry.state is ConfigEntryState.SETUP_RETRY + assert mock_setup.api.auth.call_count == 2 + 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_api.auth.call_count == 1 - assert mock_api.check_sensors.call_count == 1 + assert mock_setup.entry.state is ConfigEntryState.SETUP_RETRY + assert mock_setup.api.auth.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( - "homeassistant.components.broadlink.device.blk.gendevice", return_value=mock_api - ): - hass.config_entries.async_update_entry(mock_entry, title="New Name") + with patch(DEVICE_FACTORY, return_value=mock_setup.api): + hass.config_entries.async_update_entry(mock_setup.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") diff --git a/tests/components/broadlink/test_heartbeat.py b/tests/components/broadlink/test_heartbeat.py index 8e52a562425..de47a16c0b9 100644 --- a/tests/components/broadlink/test_heartbeat.py +++ b/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: diff --git a/tests/components/broadlink/test_remote.py b/tests/components/broadlink/test_remote.py index 2d21b588c33..abc500479ea 100644 --- a/tests/components/broadlink/test_remote.py +++ b/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_api.send_data.call_args == call(b64decode(IR_PACKET)) - assert mock_api.auth.call_count == 1 + assert mock_setup.api.send_data.call_count == 1 + assert mock_setup.api.send_data.call_args == call(b64decode(IR_PACKET)) + 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_api.send_data.call_args == call(b64decode(IR_PACKET)) - assert mock_api.auth.call_count == 1 + assert mock_setup.api.send_data.call_count == 1 + assert mock_setup.api.send_data.call_args == call(b64decode(IR_PACKET)) + assert mock_setup.api.auth.call_count == 1 diff --git a/tests/components/broadlink/test_sensors.py b/tests/components/broadlink/test_sensors.py index 5cc75c28a73..1f8f913cfe4 100644 --- a/tests/components/broadlink/test_sensors.py +++ b/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