Use entity & device registry fixtures in unifi tests (#110889)

This commit is contained in:
Jan-Philipp Benecke 2024-02-18 23:47:10 +01:00 committed by GitHub
parent 72ccc40c43
commit 22e62f42e6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 85 additions and 61 deletions

View file

@ -108,9 +108,8 @@ def mock_discovery():
@pytest.fixture
def mock_device_registry(hass):
def mock_device_registry(hass, device_registry: dr.DeviceRegistry):
"""Mock device registry."""
dev_reg = dr.async_get(hass)
config_entry = MockConfigEntry(domain="something_else")
config_entry.add_to_hass(hass)
@ -126,7 +125,7 @@ def mock_device_registry(hass):
"00:00:00:00:02:02",
)
):
dev_reg.async_get_or_create(
device_registry.async_get_or_create(
name=f"Device {idx}",
config_entry_id=config_entry.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, device)},

View file

@ -12,7 +12,10 @@ from tests.test_util.aiohttp import AiohttpClientMocker
async def test_restart_device_button(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
websocket_mock,
) -> None:
"""Test restarting device button."""
config_entry = await setup_unifi_integration(
@ -37,8 +40,7 @@ async def test_restart_device_button(
assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 1
ent_reg = er.async_get(hass)
ent_reg_entry = ent_reg.async_get("button.switch_restart")
ent_reg_entry = entity_registry.async_get("button.switch_restart")
assert ent_reg_entry.unique_id == "device_restart-00:00:00:00:01:01"
assert ent_reg_entry.entity_category is EntityCategory.CONFIG
@ -78,7 +80,10 @@ async def test_restart_device_button(
async def test_power_cycle_poe(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
websocket_mock,
) -> None:
"""Test restarting device button."""
config_entry = await setup_unifi_integration(
@ -119,8 +124,7 @@ async def test_power_cycle_poe(
assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 2
ent_reg = er.async_get(hass)
ent_reg_entry = ent_reg.async_get("button.switch_port_1_power_cycle")
ent_reg_entry = entity_registry.async_get("button.switch_port_1_power_cycle")
assert ent_reg_entry.unique_id == "power_cycle-00:00:00:00:01:01_1"
assert ent_reg_entry.entity_category is EntityCategory.CONFIG

View file

@ -242,7 +242,9 @@ async def setup_unifi_integration(
async def test_controller_setup(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
aioclient_mock: AiohttpClientMocker,
) -> None:
"""Successful setup."""
with patch(
@ -278,7 +280,7 @@ async def test_controller_setup(
assert controller.signal_options_update == "unifi-options-1"
assert controller.signal_heartbeat_missed == "unifi-heartbeat-missed"
device_entry = dr.async_get(hass).async_get_or_create(
device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id,
identifiers={(UNIFI_DOMAIN, config_entry.unique_id)},
)

View file

@ -906,7 +906,10 @@ async def test_option_ignore_wired_bug(
async def test_restoring_client(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_device_registry
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
mock_device_registry,
) -> None:
"""Verify clients are restored from clients_all if they ever was registered to entity registry."""
client = {
@ -939,15 +942,14 @@ async def test_restoring_client(
entry_id="1",
)
registry = er.async_get(hass)
registry.async_get_or_create( # Unique ID updated
entity_registry.async_get_or_create( # Unique ID updated
TRACKER_DOMAIN,
UNIFI_DOMAIN,
f'{restored["mac"]}-site_id',
suggested_object_id=restored["hostname"],
config_entry=config_entry,
)
registry.async_get_or_create( # Unique ID already updated
entity_registry.async_get_or_create( # Unique ID already updated
TRACKER_DOMAIN,
UNIFI_DOMAIN,
f'site_id-{client["mac"]}',

View file

@ -60,6 +60,7 @@ WLAN = {
async def test_wlan_qr_code(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
hass_client: ClientSessionGenerator,
snapshot: SnapshotAssertion,
@ -70,14 +71,15 @@ async def test_wlan_qr_code(
await setup_unifi_integration(hass, aioclient_mock, wlans_response=[WLAN])
assert len(hass.states.async_entity_ids(IMAGE_DOMAIN)) == 0
ent_reg = er.async_get(hass)
ent_reg_entry = ent_reg.async_get("image.ssid_1_qr_code")
ent_reg_entry = entity_registry.async_get("image.ssid_1_qr_code")
assert ent_reg_entry.unique_id == "qr_code-012345678910111213141516"
assert ent_reg_entry.disabled_by == RegistryEntryDisabler.INTEGRATION
assert ent_reg_entry.entity_category is EntityCategory.DIAGNOSTIC
# Enable entity
ent_reg.async_update_entity(entity_id="image.ssid_1_qr_code", disabled_by=None)
entity_registry.async_update_entity(
entity_id="image.ssid_1_qr_code", disabled_by=None
)
await hass.async_block_till_done()
async_fire_time_changed(

View file

@ -457,6 +457,7 @@ async def test_bandwidth_sensors(
)
async def test_uptime_sensors(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
freezer: FrozenDateTimeFactory,
mock_unifi_websocket,
@ -492,9 +493,8 @@ async def test_uptime_sensors(
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 1
assert hass.states.get("sensor.client1_uptime").state == "2021-01-01T01:00:00+00:00"
ent_reg = er.async_get(hass)
assert (
ent_reg.async_get("sensor.client1_uptime").entity_category
entity_registry.async_get("sensor.client1_uptime").entity_category
is EntityCategory.DIAGNOSTIC
)
@ -606,6 +606,7 @@ async def test_remove_sensors(
async def test_poe_port_switches(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
websocket_mock,
@ -614,13 +615,12 @@ async def test_poe_port_switches(
await setup_unifi_integration(hass, aioclient_mock, devices_response=[DEVICE_1])
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 2
ent_reg = er.async_get(hass)
ent_reg_entry = ent_reg.async_get("sensor.mock_name_port_1_poe_power")
ent_reg_entry = entity_registry.async_get("sensor.mock_name_port_1_poe_power")
assert ent_reg_entry.disabled_by == RegistryEntryDisabler.INTEGRATION
assert ent_reg_entry.entity_category is EntityCategory.DIAGNOSTIC
# Enable entity
ent_reg.async_update_entity(
entity_registry.async_update_entity(
entity_id="sensor.mock_name_port_1_poe_power", disabled_by=None
)
await hass.async_block_till_done()
@ -681,6 +681,7 @@ async def test_poe_port_switches(
async def test_wlan_client_sensors(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
websocket_mock,
@ -716,8 +717,7 @@ async def test_wlan_client_sensors(
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 1
ent_reg = er.async_get(hass)
ent_reg_entry = ent_reg.async_get("sensor.ssid_1")
ent_reg_entry = entity_registry.async_get("sensor.ssid_1")
assert ent_reg_entry.unique_id == "wlan_clients-012345678910111213141516"
assert ent_reg_entry.entity_category is EntityCategory.DIAGNOSTIC
@ -824,6 +824,7 @@ async def test_wlan_client_sensors(
)
async def test_outlet_power_readings(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
entity_id: str,
@ -838,8 +839,7 @@ async def test_outlet_power_readings(
assert len(hass.states.async_all()) == 11
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 5
ent_reg = er.async_get(hass)
ent_reg_entry = ent_reg.async_get(f"sensor.{entity_id}")
ent_reg_entry = entity_registry.async_get(f"sensor.{entity_id}")
assert ent_reg_entry.unique_id == expected_unique_id
assert ent_reg_entry.entity_category is EntityCategory.DIAGNOSTIC
@ -859,7 +859,10 @@ async def test_outlet_power_readings(
async def test_device_uptime(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
) -> None:
"""Verify that uptime sensors are working as expected."""
device = {
@ -887,9 +890,8 @@ async def test_device_uptime(
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 2
assert hass.states.get("sensor.device_uptime").state == "2021-01-01T01:00:00+00:00"
ent_reg = er.async_get(hass)
assert (
ent_reg.async_get("sensor.device_uptime").entity_category
entity_registry.async_get("sensor.device_uptime").entity_category
is EntityCategory.DIAGNOSTIC
)
@ -915,7 +917,10 @@ async def test_device_uptime(
async def test_device_temperature(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
) -> None:
"""Verify that temperature sensors are working as expected."""
device = {
@ -943,9 +948,8 @@ async def test_device_temperature(
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 3
assert hass.states.get("sensor.device_temperature").state == "30"
ent_reg = er.async_get(hass)
assert (
ent_reg.async_get("sensor.device_temperature").entity_category
entity_registry.async_get("sensor.device_temperature").entity_category
is EntityCategory.DIAGNOSTIC
)
@ -956,7 +960,10 @@ async def test_device_temperature(
async def test_device_state(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
) -> None:
"""Verify that state sensors are working as expected."""
device = {
@ -983,9 +990,8 @@ async def test_device_state(
await setup_unifi_integration(hass, aioclient_mock, devices_response=[device])
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 3
ent_reg = er.async_get(hass)
assert (
ent_reg.async_get("sensor.device_state").entity_category
entity_registry.async_get("sensor.device_state").entity_category
is EntityCategory.DIAGNOSTIC
)

View file

@ -52,7 +52,9 @@ async def test_service_setup_and_unload_not_called_if_multiple_integrations_dete
async def test_reconnect_client(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
aioclient_mock: AiohttpClientMocker,
) -> None:
"""Verify call to reconnect client is performed as expected."""
clients = [
@ -71,7 +73,6 @@ async def test_reconnect_client(
f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr",
)
device_registry = dr.async_get(hass)
device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, clients[0]["mac"])},
@ -104,14 +105,15 @@ async def test_reconnect_non_existant_device(
async def test_reconnect_device_without_mac(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
aioclient_mock: AiohttpClientMocker,
) -> None:
"""Verify no call is made if device does not have a known mac."""
config_entry = await setup_unifi_integration(hass, aioclient_mock)
aioclient_mock.clear_requests()
device_registry = dr.async_get(hass)
device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id,
connections={("other connection", "not mac")},
@ -127,7 +129,9 @@ async def test_reconnect_device_without_mac(
async def test_reconnect_client_controller_unavailable(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
aioclient_mock: AiohttpClientMocker,
) -> None:
"""Verify no call is made if controller is unavailable."""
clients = [
@ -147,7 +151,6 @@ async def test_reconnect_client_controller_unavailable(
f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr",
)
device_registry = dr.async_get(hass)
device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, clients[0]["mac"])},
@ -163,14 +166,15 @@ async def test_reconnect_client_controller_unavailable(
async def test_reconnect_client_unknown_mac(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
aioclient_mock: AiohttpClientMocker,
) -> None:
"""Verify no call is made if trying to reconnect a mac unknown to controller."""
config_entry = await setup_unifi_integration(hass, aioclient_mock)
aioclient_mock.clear_requests()
device_registry = dr.async_get(hass)
device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, "mac unknown to controller")},
@ -186,7 +190,9 @@ async def test_reconnect_client_unknown_mac(
async def test_reconnect_wired_client(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
aioclient_mock: AiohttpClientMocker,
) -> None:
"""Verify no call is made if client is wired."""
clients = [
@ -201,7 +207,6 @@ async def test_reconnect_wired_client(
aioclient_mock.clear_requests()
device_registry = dr.async_get(hass)
device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, clients[0]["mac"])},

View file

@ -816,7 +816,9 @@ async def test_not_admin(
async def test_switches(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
) -> None:
"""Test the update_items function with some clients."""
config_entry = await setup_unifi_integration(
@ -852,9 +854,10 @@ async def test_switches(
assert dpi_switch.state == "on"
assert dpi_switch.attributes["icon"] == "mdi:network"
ent_reg = er.async_get(hass)
for entry_id in ("switch.block_client_1", "switch.block_media_streaming"):
assert ent_reg.async_get(entry_id).entity_category is EntityCategory.CONFIG
assert (
entity_registry.async_get(entry_id).entity_category is EntityCategory.CONFIG
)
# Block and unblock client
aioclient_mock.clear_requests()
@ -1326,6 +1329,7 @@ async def test_option_remove_switches(
async def test_poe_port_switches(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
websocket_mock,
@ -1338,16 +1342,15 @@ async def test_poe_port_switches(
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
ent_reg = er.async_get(hass)
ent_reg_entry = ent_reg.async_get("switch.mock_name_port_1_poe")
ent_reg_entry = entity_registry.async_get("switch.mock_name_port_1_poe")
assert ent_reg_entry.disabled_by == RegistryEntryDisabler.INTEGRATION
assert ent_reg_entry.entity_category is EntityCategory.CONFIG
# Enable entity
ent_reg.async_update_entity(
entity_registry.async_update_entity(
entity_id="switch.mock_name_port_1_poe", disabled_by=None
)
ent_reg.async_update_entity(
entity_registry.async_update_entity(
entity_id="switch.mock_name_port_2_poe", disabled_by=None
)
await hass.async_block_till_done()
@ -1438,6 +1441,7 @@ async def test_poe_port_switches(
async def test_wlan_switches(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
websocket_mock,
@ -1450,8 +1454,7 @@ async def test_wlan_switches(
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
ent_reg = er.async_get(hass)
ent_reg_entry = ent_reg.async_get("switch.ssid_1")
ent_reg_entry = entity_registry.async_get("switch.ssid_1")
assert ent_reg_entry.unique_id == "wlan-012345678910111213141516"
assert ent_reg_entry.entity_category is EntityCategory.CONFIG
@ -1507,6 +1510,7 @@ async def test_wlan_switches(
async def test_port_forwarding_switches(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
websocket_mock,
@ -1531,8 +1535,7 @@ async def test_port_forwarding_switches(
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
ent_reg = er.async_get(hass)
ent_reg_entry = ent_reg.async_get("switch.unifi_network_plex")
ent_reg_entry = entity_registry.async_get("switch.unifi_network_plex")
assert ent_reg_entry.unique_id == "port_forward-5a32aa4ee4b0412345678911"
assert ent_reg_entry.entity_category is EntityCategory.CONFIG
@ -1594,7 +1597,9 @@ async def test_port_forwarding_switches(
async def test_updating_unique_id(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
) -> None:
"""Verify outlet control and poe control unique ID update works."""
poe_device = {
@ -1637,15 +1642,14 @@ async def test_updating_unique_id(
entry_id="1",
)
registry = er.async_get(hass)
registry.async_get_or_create(
entity_registry.async_get_or_create(
SWITCH_DOMAIN,
UNIFI_DOMAIN,
f'{poe_device["mac"]}-poe-1',
suggested_object_id="switch_port_1_poe",
config_entry=config_entry,
)
registry.async_get_or_create(
entity_registry.async_get_or_create(
SWITCH_DOMAIN,
UNIFI_DOMAIN,
f'{OUTLET_UP1["mac"]}-outlet-1',