Migrate calls to async_get_device in tests (part 6) (#178339)

tests/components/vodafone_station/test_coordinator.py

@@ -66,7 +66,10 @@ async def test_coordinator_device_cleanup(
     assert f"Skipping entity {DEVICE_2_HOST}" in caplog.text
 
     assert (
-        device_registry.async_get_device(identifiers={(DOMAIN, DEVICE_1_MAC)}) is None
+        device_registry.async_get_device_by_identifier(
+            (DOMAIN, DEVICE_1_MAC), mock_config_entry.entry_id
+        )
+        is None
     )
     assert f"Removing device: {DEVICE_1_HOST}" in caplog.text
 

tests/components/voip/test_devices.py

@@ -20,13 +20,14 @@ async def test_device_registry_info(
     voip_devices: VoIPDevices,
     call_info: CallInfo,
     device_registry: dr.DeviceRegistry,
+    config_entry: MockConfigEntry,
 ) -> None:
     """Test info in device registry."""
     voip_device = voip_devices.async_get_or_create(call_info)
     assert not voip_device.async_allow_call(hass)
 
-    device = device_registry.async_get_device(
-        identifiers={(DOMAIN, call_info.caller_endpoint.uri)}
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, call_info.caller_endpoint.uri), config_entry.entry_id
     )
     assert device is not None
     assert device.name == call_info.caller_endpoint.host
@@ -40,8 +41,8 @@ async def test_device_registry_info(
 
     assert not voip_device.async_allow_call(hass)
 
-    device = device_registry.async_get_device(
-        identifiers={(DOMAIN, call_info.caller_endpoint.uri)}
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, call_info.caller_endpoint.uri), config_entry.entry_id
     )
     assert device.sw_version == "2.0.0.0"
 
@@ -51,14 +52,15 @@ async def test_device_registry_info_from_unknown_phone(
     voip_devices: VoIPDevices,
     call_info: CallInfo,
     device_registry: dr.DeviceRegistry,
+    config_entry: MockConfigEntry,
 ) -> None:
     """Test info in device registry from unknown phone."""
     call_info.headers["user-agent"] = "Unknown"
     voip_device = voip_devices.async_get_or_create(call_info)
     assert not voip_device.async_allow_call(hass)
 
-    device = device_registry.async_get_device(
-        identifiers={(DOMAIN, call_info.caller_endpoint.uri)}
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, call_info.caller_endpoint.uri), config_entry.entry_id
     )
     assert device.manufacturer is None
     assert device.model == "Unknown"
@@ -70,13 +72,14 @@ async def test_device_registry_info_update_contact(
     voip_devices: VoIPDevices,
     call_info: CallInfo,
     device_registry: dr.DeviceRegistry,
+    config_entry: MockConfigEntry,
 ) -> None:
     """Test info in device registry."""
     voip_device = voip_devices.async_get_or_create(call_info)
     assert not voip_device.async_allow_call(hass)
 
-    device = device_registry.async_get_device(
-        identifiers={(DOMAIN, call_info.caller_endpoint.uri)}
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, call_info.caller_endpoint.uri), config_entry.entry_id
     )
     assert device is not None
     assert device.name == call_info.caller_endpoint.host
@@ -92,8 +95,8 @@ async def test_device_registry_info_update_contact(
     assert voip_device.contact == SipEndpoint("Test <sip:example.com:5061>")
     assert not voip_device.async_allow_call(hass)
 
-    device = device_registry.async_get_device(
-        identifiers={(DOMAIN, call_info.caller_endpoint.uri)}
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, call_info.caller_endpoint.uri), config_entry.entry_id
     )
     assert device.sw_version == "2.0.0.0"
 
@@ -179,13 +182,16 @@ async def test_device_registry_migration(
     call_info: CallInfo,
     entity_registry: er.EntityRegistry,
     device_registry: dr.DeviceRegistry,
+    config_entry: MockConfigEntry,
 ) -> None:
     """Test info in device registry migrates old devices."""
     voip_device = voip_devices.async_get_or_create(call_info)
     new_id = call_info.caller_endpoint.uri
     assert voip_device.voip_id == new_id
 
-    device = device_registry.async_get_device(identifiers={(DOMAIN, new_id)})
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, new_id), config_entry.entry_id
+    )
     assert device is not None
     assert device.id == legacy_dev_reg_entry.id
     assert device.identifiers == {(DOMAIN, new_id)}

tests/components/waqi/test_init.py

@@ -111,7 +111,11 @@ async def test_migration_from_v1(
     assert entity.config_subentry_id == subentry.subentry_id
     assert entity.config_entry_id == entry.entry_id
 
-    assert (device := device_registry.async_get_device(identifiers={(DOMAIN, "4584")}))
+    assert (
+        device := device_registry.async_get_device_by_identifier(
+            (DOMAIN, "4584"), mock_config_entry.entry_id
+        )
+    )
     assert device.identifiers == {(DOMAIN, "4584")}
     assert device.id == device_1.id
     assert device.config_entries == {mock_config_entry.entry_id}
@@ -129,7 +133,11 @@ async def test_migration_from_v1(
     assert entity.unique_id == "4585_air_quality"
     assert entity.config_subentry_id == subentry.subentry_id
     assert entity.config_entry_id == entry.entry_id
-    assert (device := device_registry.async_get_device(identifiers={(DOMAIN, "4585")}))
+    assert (
+        device := device_registry.async_get_device_by_identifier(
+            (DOMAIN, "4585"), mock_config_entry.entry_id
+        )
+    )
     assert device.identifiers == {(DOMAIN, "4585")}
     assert device.id == device_2.id
     assert device.config_entries == {mock_config_entry.entry_id}
@@ -315,11 +323,11 @@ async def test_migration_from_v1_disabled(
         assert subentry.data == {CONF_STATION_NUMBER: int(subentry.unique_id)}
         assert "de Jongweg" in subentry.title
 
-    assert not device_registry.async_get_device(
-        identifiers={(DOMAIN, mock_config_entry.entry_id)}
+    assert not device_registry.async_get_device_by_identifier(
+        (DOMAIN, mock_config_entry.entry_id), mock_config_entry.entry_id
     )
-    assert not device_registry.async_get_device(
-        identifiers={(DOMAIN, mock_config_entry_2.entry_id)}
+    assert not device_registry.async_get_device_by_identifier(
+        (DOMAIN, mock_config_entry_2.entry_id), mock_config_entry_2.entry_id
     )
 
     for idx, subentry in enumerate(station_subentries):
@@ -331,8 +339,9 @@ async def test_migration_from_v1_disabled(
         assert entity.disabled_by is subentry_data["entity_disabled_by"]
 
         assert (
-            device := device_registry.async_get_device(
-                identifiers={(DOMAIN, subentry.unique_id)}
+            device := device_registry.async_get_device_by_identifier(
+                (DOMAIN, subentry.unique_id),
+                mock_config_entries[main_config_entry].entry_id,
             )
         )
         assert device.identifiers == {(DOMAIN, subentry.unique_id)}

tests/components/watts/test_init.py

@@ -203,10 +203,16 @@ async def test_dynamic_device_creation(
     """Test new devices are created dynamically."""
     await setup_integration(hass, mock_config_entry)
 
-    assert device_registry.async_get_device(identifiers={(DOMAIN, "thermostat_123")})
-    assert device_registry.async_get_device(identifiers={(DOMAIN, "thermostat_456")})
+    assert device_registry.async_get_device_by_identifier(
+        (DOMAIN, "thermostat_123"), mock_config_entry.entry_id
+    )
+    assert device_registry.async_get_device_by_identifier(
+        (DOMAIN, "thermostat_456"), mock_config_entry.entry_id
+    )
     assert (
-        device_registry.async_get_device(identifiers={(DOMAIN, "thermostat_789")})
+        device_registry.async_get_device_by_identifier(
+            (DOMAIN, "thermostat_789"), mock_config_entry.entry_id
+        )
         is None
     )
 
@@ -234,8 +240,8 @@ async def test_dynamic_device_creation(
     async_fire_time_changed(hass)
     await hass.async_block_till_done()
 
-    new_device_entry = device_registry.async_get_device(
-        identifiers={(DOMAIN, "thermostat_789")}
+    new_device_entry = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "thermostat_789"), mock_config_entry.entry_id
     )
     assert new_device_entry is not None
     assert new_device_entry.name == "Kitchen Thermostat"
@@ -254,11 +260,11 @@ async def test_stale_device_removal(
     """Test stale devices are removed dynamically."""
     await setup_integration(hass, mock_config_entry)
 
-    device_123 = device_registry.async_get_device(
-        identifiers={(DOMAIN, "thermostat_123")}
+    device_123 = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "thermostat_123"), mock_config_entry.entry_id
     )
-    device_456 = device_registry.async_get_device(
-        identifiers={(DOMAIN, "thermostat_456")}
+    device_456 = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "thermostat_456"), mock_config_entry.entry_id
     )
     assert device_123 is not None
     assert device_456 is not None
@@ -275,8 +281,8 @@ async def test_stale_device_removal(
     await hass.async_block_till_done()
 
     # Verify thermostat_456 has been removed
-    device_456_after_removal = device_registry.async_get_device(
-        identifiers={(DOMAIN, "thermostat_456")}
+    device_456_after_removal = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "thermostat_456"), mock_config_entry.entry_id
     )
     assert device_456_after_removal is None
 

tests/components/webostv/test_device_trigger.py

@@ -24,9 +24,11 @@ async def test_get_triggers(
     hass: HomeAssistant, device_registry: dr.DeviceRegistry, client
 ) -> None:
     """Test we get the expected triggers."""
-    await setup_webostv(hass)
+    entry = await setup_webostv(hass)
 
-    device = device_registry.async_get_device(identifiers={(DOMAIN, FAKE_UUID)})
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, FAKE_UUID), entry.entry_id
+    )
 
     turn_on_trigger = {
         "platform": "device",
@@ -49,9 +51,11 @@ async def test_if_fires_on_turn_on_request(
     client,
 ) -> None:
     """Test for turn_on and turn_off triggers firing."""
-    await setup_webostv(hass)
+    entry = await setup_webostv(hass)
 
-    device = device_registry.async_get_device(identifiers={(DOMAIN, FAKE_UUID)})
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, FAKE_UUID), entry.entry_id
+    )
 
     assert await async_setup_component(
         hass,

tests/components/webostv/test_media_player.py

@@ -324,7 +324,9 @@ async def test_device_info_startup_off(
 
     assert hass.states.get(ENTITY_ID).state == STATE_OFF
 
-    device = device_registry.async_get_device(identifiers={(DOMAIN, entry.unique_id)})
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, entry.unique_id), entry.entry_id
+    )
 
     assert device
     assert device.identifiers == {(DOMAIN, entry.unique_id)}
@@ -363,7 +365,9 @@ async def test_entity_attributes(
     assert attrs[ATTR_MEDIA_TITLE] == "Channel Name 2"
 
     # Device Info
-    device = device_registry.async_get_device(identifiers={(DOMAIN, entry.unique_id)})
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, entry.unique_id), entry.entry_id
+    )
     assert device == snapshot
 
     # Sound output when off

tests/components/webostv/test_trigger.py

@@ -25,9 +25,11 @@ async def test_webostv_turn_on_trigger_device_id(
     client,
 ) -> None:
     """Test for turn_on triggers by device_id firing."""
-    await setup_webostv(hass)
+    entry = await setup_webostv(hass)
 
-    device = device_registry.async_get_device(identifiers={(DOMAIN, FAKE_UUID)})
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, FAKE_UUID), entry.entry_id
+    )
 
     assert await async_setup_component(
         hass,

tests/components/wiim/test_entity.py

@@ -25,8 +25,8 @@ async def test_device_info_uses_http_api_url(
 
     await setup_integration(hass, mock_config_entry)
 
-    device_entry = device_registry.async_get_device(
-        identifiers={(DOMAIN, mock_wiim_device.udn)}
+    device_entry = device_registry.async_get_device_by_identifier(
+        (DOMAIN, mock_wiim_device.udn), mock_config_entry.entry_id
     )
     assert device_entry is not None
     assert device_entry.configuration_url == mock_wiim_device.http_api_url

tests/components/withings/test_init.py

@@ -710,7 +710,9 @@ async def test_devices(
     await hass.async_block_till_done()
 
     for device_id in ("12345", "f998be4b9ccc9e136fd8cd8e8e344c31ec3b271d"):
-        device = device_registry.async_get_device({(DOMAIN, device_id)})
+        device = device_registry.async_get_device_by_identifier(
+            (DOMAIN, device_id), webhook_config_entry.entry_id
+        )
         assert device is not None
         assert device == snapshot(name=device_id)
 

tests/components/withings/test_sensor.py

@@ -397,8 +397,9 @@ async def test_device_sensors_created_when_device_data_received(
     await hass.async_block_till_done()
 
     assert hass.states.get("sensor.body_battery")
-    assert device_registry.async_get_device(
-        {(DOMAIN, "f998be4b9ccc9e136fd8cd8e8e344c31ec3b271d")}
+    assert device_registry.async_get_device_by_identifier(
+        (DOMAIN, "f998be4b9ccc9e136fd8cd8e8e344c31ec3b271d"),
+        polling_config_entry.entry_id,
     )
 
     withings.get_devices.return_value = []
@@ -408,8 +409,9 @@ async def test_device_sensors_created_when_device_data_received(
     await hass.async_block_till_done()
 
     assert hass.states.get("sensor.body_battery") is None
-    assert not device_registry.async_get_device(
-        {(DOMAIN, "f998be4b9ccc9e136fd8cd8e8e344c31ec3b271d")}
+    assert not device_registry.async_get_device_by_identifier(
+        (DOMAIN, "f998be4b9ccc9e136fd8cd8e8e344c31ec3b271d"),
+        polling_config_entry.entry_id,
     )
 
 

tests/components/wmspro/test_cover.py

@@ -56,7 +56,9 @@ async def test_cover_device(
     assert len(mock_hub_configuration.mock_calls) == 1
     assert len(mock_hub_status.mock_calls) == len(mock_hub_configuration.destinations)
 
-    device_entry = device_registry.async_get_device(identifiers={(DOMAIN, "58717")})
+    device_entry = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "58717"), mock_config_entry.entry_id
+    )
     assert device_entry is not None
     assert device_entry == snapshot
 

tests/components/wmspro/test_light.py

@@ -44,7 +44,9 @@ async def test_light_device(
     assert len(mock_hub_configuration.mock_calls) == 1
     assert len(mock_hub_status.mock_calls) == len(mock_hub_configuration.destinations)
 
-    device_entry = device_registry.async_get_device(identifiers={(DOMAIN, "97358")})
+    device_entry = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "97358"), mock_config_entry.entry_id
+    )
     assert device_entry is not None
     assert device_entry == snapshot
 

tests/components/wmspro/test_scene.py

@@ -36,7 +36,9 @@ async def test_scene_room_device(
     assert len(mock_hub_configuration.mock_calls) == 1
     assert len(mock_dest_refresh.mock_calls) == 2
 
-    device_entry = device_registry.async_get_device(identifiers={(DOMAIN, "42581")})
+    device_entry = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "42581"), mock_config_entry.entry_id
+    )
     assert device_entry is not None
     assert device_entry == snapshot
 

tests/components/wmspro/test_switch.py

@@ -44,7 +44,9 @@ async def test_switch_device(
     assert len(mock_hub_configuration.mock_calls) == 1
     assert len(mock_hub_status.mock_calls) == len(mock_hub_configuration.destinations)
 
-    device_entry = device_registry.async_get_device(identifiers={(DOMAIN, "499120")})
+    device_entry = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "499120"), mock_config_entry.entry_id
+    )
     assert device_entry is not None
     assert device_entry == snapshot
 

tests/components/wolflink/test_init.py

@@ -393,8 +393,12 @@ async def test_setup_multiple_devices(
     await setup_integration(hass, mock_config_entry)
 
     # Both devices must appear in the device registry.
-    first = device_registry.async_get_device({(DOMAIN, "1234")})
-    second = device_registry.async_get_device({(DOMAIN, "9999")})
+    first = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "1234"), mock_config_entry.entry_id
+    )
+    second = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "9999"), mock_config_entry.entry_id
+    )
     assert first is not None
     assert second is not None
     assert first.name == "first-device"

tests/components/wolflink/test_sensor.py

@@ -24,7 +24,9 @@ async def test_device_entry(
     await hass.config_entries.async_setup(mock_config_entry.entry_id)
     await hass.async_block_till_done()
 
-    device = device_registry.async_get_device({(mock_config_entry.domain, "1234")})
+    device = device_registry.async_get_device_by_identifier(
+        (mock_config_entry.domain, "1234"), mock_config_entry.entry_id
+    )
     assert device == snapshot
 
 

tests/components/wyoming/test_devices.py

@@ -20,8 +20,8 @@ async def test_device_registry_info(
 
     # Satellite uses config entry id since only one satellite per entry is
     # supported.
-    device = device_registry.async_get_device(
-        identifiers={(DOMAIN, satellite_config_entry.entry_id)}
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, satellite_config_entry.entry_id), satellite_config_entry.entry_id
     )
     assert device is not None
     assert device.name == "Test Satellite"

tests/components/xbox/test_init.py

@@ -219,8 +219,18 @@ async def test_dynamic_devices(
 
     assert config_entry.state is ConfigEntryState.LOADED
 
-    assert device_registry.async_get_device({(DOMAIN, "ABCDEFG")}) is None
-    assert device_registry.async_get_device({(DOMAIN, "HIJKLMN")}) is None
+    assert (
+        device_registry.async_get_device_by_identifier(
+            (DOMAIN, "ABCDEFG"), config_entry.entry_id
+        )
+        is None
+    )
+    assert (
+        device_registry.async_get_device_by_identifier(
+            (DOMAIN, "HIJKLMN"), config_entry.entry_id
+        )
+        is None
+    )
 
     xbox_live_client.smartglass.get_console_list.return_value = SmartglassConsoleList(
         **await async_load_json_object_fixture(
@@ -232,8 +242,12 @@ async def test_dynamic_devices(
     async_fire_time_changed(hass)
     await hass.async_block_till_done()
 
-    assert device_registry.async_get_device({(DOMAIN, "ABCDEFG")})
-    assert device_registry.async_get_device({(DOMAIN, "HIJKLMN")})
+    assert device_registry.async_get_device_by_identifier(
+        (DOMAIN, "ABCDEFG"), config_entry.entry_id
+    )
+    assert device_registry.async_get_device_by_identifier(
+        (DOMAIN, "HIJKLMN"), config_entry.entry_id
+    )
 
     xbox_live_client.smartglass.get_console_list.return_value = SmartglassConsoleList(
         **await async_load_json_object_fixture(
@@ -245,5 +259,15 @@ async def test_dynamic_devices(
     async_fire_time_changed(hass)
     await hass.async_block_till_done()
 
-    assert device_registry.async_get_device({(DOMAIN, "ABCDEFG")}) is None
-    assert device_registry.async_get_device({(DOMAIN, "HIJKLMN")}) is None
+    assert (
+        device_registry.async_get_device_by_identifier(
+            (DOMAIN, "ABCDEFG"), config_entry.entry_id
+        )
+        is None
+    )
+    assert (
+        device_registry.async_get_device_by_identifier(
+            (DOMAIN, "HIJKLMN"), config_entry.entry_id
+        )
+        is None
+    )

tests/components/xiaomi_ble/test_device_trigger.py

@@ -190,7 +190,9 @@ async def test_get_triggers_button(
     await hass.async_block_till_done()
     assert len(events) == 1
 
-    device = device_registry.async_get_device(identifiers={get_device_id(mac)})
+    device = device_registry.async_get_device_by_identifier(
+        get_device_id(mac), entry.entry_id
+    )
     assert device
     expected_trigger = {
         CONF_PLATFORM: "device",
@@ -231,7 +233,9 @@ async def test_get_triggers_double_button(
     await hass.async_block_till_done()
     assert len(events) == 1
 
-    device = device_registry.async_get_device(identifiers={get_device_id(mac)})
+    device = device_registry.async_get_device_by_identifier(
+        get_device_id(mac), entry.entry_id
+    )
     assert device
     expected_trigger = {
         CONF_PLATFORM: "device",
@@ -273,7 +277,9 @@ async def test_get_triggers_lock(
     await hass.async_block_till_done()
     assert len(events) == 1
 
-    device = device_registry.async_get_device(identifiers={get_device_id(mac)})
+    device = device_registry.async_get_device_by_identifier(
+        get_device_id(mac), entry.entry_id
+    )
     assert device
     expected_trigger = {
         CONF_PLATFORM: "device",
@@ -310,7 +316,9 @@ async def test_get_triggers_motion(
     await hass.async_block_till_done()
     assert len(events) == 1
 
-    device = device_registry.async_get_device(identifiers={get_device_id(mac)})
+    device = device_registry.async_get_device_by_identifier(
+        get_device_id(mac), entry.entry_id
+    )
     assert device
     expected_trigger = {
         CONF_PLATFORM: "device",
@@ -413,7 +421,9 @@ async def test_if_fires_on_button_press(
     # wait for the device being created
     await hass.async_block_till_done()
 
-    device = device_registry.async_get_device(identifiers={get_device_id(mac)})
+    device = device_registry.async_get_device_by_identifier(
+        get_device_id(mac), entry.entry_id
+    )
     device_id = device.id
 
     assert await async_setup_component(
@@ -476,7 +486,9 @@ async def test_if_fires_on_double_button_long_press(
     # wait for the device being created
     await hass.async_block_till_done()
 
-    device = device_registry.async_get_device(identifiers={get_device_id(mac)})
+    device = device_registry.async_get_device_by_identifier(
+        get_device_id(mac), entry.entry_id
+    )
     device_id = device.id
 
     assert await async_setup_component(
@@ -535,7 +547,9 @@ async def test_if_fires_on_motion_detected(
     # wait for the device being created
     await hass.async_block_till_done()
 
-    device = device_registry.async_get_device(identifiers={get_device_id(mac)})
+    device = device_registry.async_get_device_by_identifier(
+        get_device_id(mac), entry.entry_id
+    )
     device_id = device.id
 
     assert await async_setup_component(
@@ -591,7 +605,9 @@ async def test_automation_with_invalid_trigger_type(
     # wait for the event
     await hass.async_block_till_done()
 
-    device = device_registry.async_get_device(identifiers={get_device_id(mac)})
+    device = device_registry.async_get_device_by_identifier(
+        get_device_id(mac), entry.entry_id
+    )
     device_id = device.id
 
     assert await async_setup_component(
@@ -640,7 +656,9 @@ async def test_automation_with_invalid_trigger_event_property(
     # wait for the event
     await hass.async_block_till_done()
 
-    device = device_registry.async_get_device(identifiers={get_device_id(mac)})
+    device = device_registry.async_get_device_by_identifier(
+        get_device_id(mac), entry.entry_id
+    )
     device_id = device.id
 
     assert await async_setup_component(

tests/components/xthings_cloud/test_init.py

@@ -24,7 +24,9 @@ async def test_devices(
     await setup_integration(hass, mock_config_entry)
 
     for device in mock_api_client.async_get_devices.return_value:
-        device_entry = device_registry.async_get_device({(DOMAIN, device["id"])})
+        device_entry = device_registry.async_get_device_by_identifier(
+            (DOMAIN, device["id"]), mock_config_entry.entry_id
+        )
 
         assert device_entry is not None
         assert device_entry == snapshot(name=device["model"])

tests/components/yale/test_binary_sensor.py

@@ -309,9 +309,11 @@ async def test_doorbell_device_registry(
 ) -> None:
     """Test creation of a lock with doorsense and bridge ands up in the registry."""
     doorbell_one = await _mock_doorbell_from_fixture(hass, "get_doorbell.offline.json")
-    await _create_yale_with_devices(hass, [doorbell_one])
+    config_entry, _ = await _create_yale_with_devices(hass, [doorbell_one])
 
-    reg_device = device_registry.async_get_device(identifiers={("yale", "tmt100")})
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("yale", "tmt100"), config_entry.entry_id
+    )
     assert reg_device == snapshot
 
 

tests/components/yale/test_lock.py

@@ -40,10 +40,10 @@ async def test_lock_device_registry(
 ) -> None:
     """Test creation of a lock with doorsense and bridge ands up in the registry."""
     lock_one = await _mock_doorsense_enabled_yale_lock_detail(hass)
-    await _create_yale_with_devices(hass, [lock_one])
+    entry, _ = await _create_yale_with_devices(hass, [lock_one])
 
-    reg_device = device_registry.async_get_device(
-        identifiers={("yale", "online_with_doorsense")}
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("yale", "online_with_doorsense"), entry.entry_id
     )
     assert reg_device == snapshot
 

tests/components/yolink/test_device_trigger.py

@@ -150,7 +150,9 @@ async def test_if_fires_on_event(
         },
     )
 
-    device = device_registry.async_get_device(connections={connection})
+    device = device_registry.async_get_device_by_connection(
+        connection, config_entry.entry_id
+    )
     assert device is not None
     # Fake remote button long press.
     hass.bus.async_fire(

tests/components/yoto/test_init.py

@@ -397,7 +397,10 @@ async def test_stale_device_removed(
     """A player removed from the account has its device dropped."""
     await setup_integration(hass, mock_config_entry)
     assert (
-        device_registry.async_get_device(identifiers={(DOMAIN, PLAYER_ID)}) is not None
+        device_registry.async_get_device_by_identifier(
+            (DOMAIN, PLAYER_ID), mock_config_entry.entry_id
+        )
+        is not None
     )
 
     mock_yoto_client.players.clear()
@@ -405,5 +408,10 @@ async def test_stale_device_removed(
     async_fire_time_changed(hass)
     await hass.async_block_till_done()
 
-    assert device_registry.async_get_device(identifiers={(DOMAIN, PLAYER_ID)}) is None
+    assert (
+        device_registry.async_get_device_by_identifier(
+            (DOMAIN, PLAYER_ID), mock_config_entry.entry_id
+        )
+        is None
+    )
     mock_yoto_client.unsubscribe_player_events.assert_called_once_with(PLAYER_ID)

tests/components/youtube/test_init.py

@@ -130,8 +130,8 @@ async def test_device_info(
 
     entry = hass.config_entries.async_entries(DOMAIN)[0]
     channel_id = entry.options[CONF_CHANNELS][0]
-    device = device_registry.async_get_device(
-        identifiers={(DOMAIN, f"{entry.entry_id}_{channel_id}")}
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, f"{entry.entry_id}_{channel_id}"), entry.entry_id
     )
 
     assert device.entry_type is dr.DeviceEntryType.SERVICE

tests/components/zeversolar/test_diagnostics.py

@@ -37,8 +37,8 @@ async def test_device_diagnostics(
     init_integration: MockConfigEntry,
 ) -> None:
     """Test device diagnostics."""
-    device = device_registry.async_get_device(
-        identifiers={(DOMAIN, MOCK_SERIAL_NUMBER)}
+    device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, MOCK_SERIAL_NUMBER), init_integration.entry_id
     )
 
     assert (

tests/components/zha/test_device_action.py

@@ -21,7 +21,11 @@ from homeassistant.setup import async_setup_component
 
 from .conftest import SIG_EP_INPUT, SIG_EP_OUTPUT, SIG_EP_PROFILE, SIG_EP_TYPE
 
-from tests.common import async_get_device_automations, async_mock_service
+from tests.common import (
+    MockConfigEntry,
+    async_get_device_automations,
+    async_mock_service,
+)
 
 SHORT_PRESS = "remote_button_short_press"
 COMMAND = "command"
@@ -52,6 +56,7 @@ async def test_get_actions(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,
     entity_registry: er.EntityRegistry,
+    config_entry: MockConfigEntry,
     setup_zha: Callable[..., Coroutine[None]],
     zigpy_device_mock: Callable[..., Device],
 ) -> None:
@@ -80,7 +85,9 @@ async def test_get_actions(
     await hass.async_block_till_done(wait_background_tasks=True)
     ieee_address = str(zigpy_device.ieee)
 
-    reg_device = device_registry.async_get_device(identifiers={(DOMAIN, ieee_address)})
+    reg_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, ieee_address), config_entry.entry_id
+    )
     siren_level_select = entity_registry.async_get(
         "select.fakemanufacturer_fakemodel_default_siren_level"
     )
@@ -138,6 +145,7 @@ async def test_get_actions(
 async def test_action(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,
+    config_entry: MockConfigEntry,
     setup_zha: Callable[..., Coroutine[None]],
     zigpy_device_mock: Callable[..., Device],
 ) -> None:
@@ -168,7 +176,9 @@ async def test_action(
     await hass.async_block_till_done(wait_background_tasks=True)
     ieee_address = str(zigpy_device.ieee)
 
-    reg_device = device_registry.async_get_device(identifiers={(DOMAIN, ieee_address)})
+    reg_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, ieee_address), config_entry.entry_id
+    )
 
     with patch(
         "zigpy.zcl.Cluster.request",

tests/components/zha/test_device_trigger.py

@@ -79,8 +79,9 @@ async def test_triggers(
     await gateway.async_device_initialized(zha_device.device)
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    reg_device = device_registry.async_get_device(
-        identifiers={("zha", str(zha_device.ieee))}
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("zha", str(zha_device.ieee)),
+        hass.config_entries.async_entries("zha")[0].entry_id,
     )
 
     triggers = await async_get_device_automations(
@@ -160,8 +161,9 @@ async def test_no_triggers(
     await gateway.async_device_initialized(zha_device.device)
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    reg_device = device_registry.async_get_device(
-        identifiers={("zha", str(zha_device.ieee))}
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("zha", str(zha_device.ieee)),
+        hass.config_entries.async_entries("zha")[0].entry_id,
     )
 
     triggers = await async_get_device_automations(
@@ -211,8 +213,9 @@ async def test_if_fires_on_event(
     await gateway.async_device_initialized(zha_device.device)
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    reg_device = device_registry.async_get_device(
-        identifiers={("zha", str(zha_device.ieee))}
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("zha", str(zha_device.ieee)),
+        hass.config_entries.async_entries("zha")[0].entry_id,
     )
 
     assert await async_setup_component(
@@ -276,8 +279,9 @@ async def test_device_offline_fires(
     await gateway.async_device_initialized(zha_device.device)
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    reg_device = device_registry.async_get_device(
-        identifiers={("zha", str(zha_device.ieee))}
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("zha", str(zha_device.ieee)),
+        hass.config_entries.async_entries("zha")[0].entry_id,
     )
 
     assert await async_setup_component(
@@ -332,8 +336,9 @@ async def test_exception_no_triggers(
     await gateway.async_device_initialized(zha_device.device)
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    reg_device = device_registry.async_get_device(
-        identifiers={("zha", str(zha_device.ieee))}
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("zha", str(zha_device.ieee)),
+        hass.config_entries.async_entries("zha")[0].entry_id,
     )
 
     await async_setup_component(
@@ -392,8 +397,9 @@ async def test_exception_bad_trigger(
     await gateway.async_device_initialized(zha_device.device)
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    reg_device = device_registry.async_get_device(
-        identifiers={("zha", str(zha_device.ieee))}
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("zha", str(zha_device.ieee)),
+        hass.config_entries.async_entries("zha")[0].entry_id,
     )
 
     await async_setup_component(
@@ -457,8 +463,8 @@ async def test_validate_trigger_config_missing_info(
     # it be pulled from the current device, making it impossible to validate triggers
     await hass.config_entries.async_unload(config_entry.entry_id)
 
-    reg_device = device_registry.async_get_device(
-        identifiers={("zha", str(zha_device.ieee))}
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("zha", str(zha_device.ieee)), config_entry.entry_id
     )
 
     assert await async_setup_component(
@@ -532,8 +538,8 @@ async def test_validate_trigger_config_unloaded_bad_info(
 
     await hass.config_entries.async_unload(config_entry.entry_id)
 
-    reg_device = device_registry.async_get_device(
-        identifiers={("zha", str(zha_device.ieee))}
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("zha", str(zha_device.ieee)), config_entry.entry_id
     )
 
     assert await async_setup_component(

tests/components/zha/test_diagnostics.py

@@ -129,8 +129,8 @@ async def test_diagnostics_for_device(
         security.IasZone.AttributeDefs.current_zone_sensitivity_level
     )
 
-    device = device_registry.async_get_device(
-        identifiers={("zha", str(zha_device_proxy.device.ieee))}
+    device = device_registry.async_get_device_by_identifier(
+        ("zha", str(zha_device_proxy.device.ieee)), config_entry.entry_id
     )
     assert device
     diagnostics_data = await get_diagnostics_for_device(

tests/components/zha/test_logbook.py

@@ -98,7 +98,9 @@ async def test_zha_logbook_event_device_with_triggers(
 
     ieee_address = str(zha_device.device.ieee)
 
-    reg_device = device_registry.async_get_device(identifiers={("zha", ieee_address)})
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("zha", ieee_address), hass.config_entries.async_entries("zha")[0].entry_id
+    )
 
     hass.config.components.add("recorder")
     assert await async_setup_component(hass, "logbook", {})
@@ -174,7 +176,9 @@ async def test_zha_logbook_event_device_no_triggers(
 
     _zigpy_device, zha_device = mock_devices
     ieee_address = str(zha_device.device.ieee)
-    reg_device = device_registry.async_get_device(identifiers={("zha", ieee_address)})
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("zha", ieee_address), hass.config_entries.async_entries("zha")[0].entry_id
+    )
 
     hass.config.components.add("recorder")
     assert await async_setup_component(hass, "logbook", {})

tests/components/zha/test_update.py

@@ -715,13 +715,15 @@ async def test_update_version_sync_device_registry(
     zha_device, _, _, _ = await setup_test_data(hass, zigpy_device_mock)
 
     zha_device.device.async_update_firmware_version("0x12345678")
-    reg_device = device_registry.async_get_device(
-        identifiers={("zha", str(zha_device.device.ieee))}
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("zha", str(zha_device.device.ieee)),
+        hass.config_entries.async_entries("zha")[0].entry_id,
     )
     assert reg_device.sw_version == "0x12345678"
 
     zha_device.device.async_update_firmware_version("0xabcd1234")
-    reg_device = device_registry.async_get_device(
-        identifiers={("zha", str(zha_device.device.ieee))}
+    reg_device = device_registry.async_get_device_by_identifier(
+        ("zha", str(zha_device.device.ieee)),
+        hass.config_entries.async_entries("zha")[0].entry_id,
     )
     assert reg_device.sw_version == "0xabcd1234"

tests/components/zwave_me/test_remove_stale_devices.py

@@ -63,13 +63,12 @@ async def test_remove_stale_devices(
         await hass.config_entries.async_setup(config_entry.entry_id)
     assert (
         bool(
-            device_registry.async_get_device(
-                identifiers={
-                    (
-                        "zwave_me",
-                        f"{config_entry.unique_id}-{identifier}",
-                    )
-                }
+            device_registry.async_get_device_by_identifier(
+                (
+                    "zwave_me",
+                    f"{config_entry.unique_id}-{identifier}",
+                ),
+                config_entry.entry_id,
             )
         )
         == should_exist

tests/helpers/test_entity.py

@@ -1630,7 +1630,9 @@ async def test_friendly_name_updated(
     state = hass.states.async_all()[0]
     assert state.attributes.get(ATTR_FRIENDLY_NAME) == expected_friendly_name1
 
-    device = device_registry.async_get_device(identifiers={("hue", "1234")})
+    device = device_registry.async_get_device_by_identifier(
+        ("hue", "1234"), config_entry.entry_id
+    )
     device_registry.async_update_device(device.id, name_by_user="Device Bla2")
     await hass.async_block_till_done()
 

tests/helpers/test_entity_platform.py

@@ -1508,13 +1508,17 @@ async def test_device_info_called(
 
     assert len(hass.states.async_entity_ids()) == 3
 
-    device = device_registry.async_get_device(identifiers={("hue", "1234")})
+    device = device_registry.async_get_device_by_identifier(
+        ("hue", "1234"), config_entry.entry_id
+    )
     assert device == snapshot
     assert device.config_entries == {config_entry.entry_id}
     assert device.config_entries_subentries == {config_entry.entry_id: {None}}
     assert device.primary_config_entry == config_entry.entry_id
     assert device.via_device_id == via.id
-    device = device_registry.async_get_device(identifiers={("hue", "efgh")})
+    device = device_registry.async_get_device_by_identifier(
+        ("hue", "efgh"), config_entry.entry_id
+    )
     assert device == snapshot
     assert device.config_entries == {config_entry.entry_id}
     assert device.config_entries_subentries == {
@@ -1568,8 +1572,8 @@ async def test_device_info_not_overrides(
     assert await entity_platform.async_setup_entry(config_entry)
     await hass.async_block_till_done()
 
-    device2 = device_registry.async_get_device(
-        connections={(dr.CONNECTION_NETWORK_MAC, "abcd")}
+    device2 = device_registry.async_get_device_by_connection(
+        (dr.CONNECTION_NETWORK_MAC, "abcd"), config_entry.entry_id
     )
     assert device2 is not None
     assert device.id == device2.id
@@ -1622,7 +1626,9 @@ async def test_device_info_homeassistant_url(
 
     assert len(hass.states.async_entity_ids()) == 1
 
-    device = device_registry.async_get_device(identifiers={("mqtt", "1234")})
+    device = device_registry.async_get_device_by_identifier(
+        ("mqtt", "1234"), config_entry.entry_id
+    )
     assert device is not None
     assert device.identifiers == {("mqtt", "1234")}
     assert device.configuration_url == "homeassistant://config/mqtt"
@@ -1674,7 +1680,9 @@ async def test_device_info_change_to_no_url(
 
     assert len(hass.states.async_entity_ids()) == 1
 
-    device = device_registry.async_get_device(identifiers={("mqtt", "1234")})
+    device = device_registry.async_get_device_by_identifier(
+        ("mqtt", "1234"), config_entry.entry_id
+    )
     assert device is not None
     assert device.identifiers == {("mqtt", "1234")}
     assert device.configuration_url is None
@@ -2762,8 +2770,8 @@ async def test_device_name_defaulting_config_entry(
     assert await entity_platform.async_setup_entry(config_entry)
     await hass.async_block_till_done()
 
-    device = device_registry.async_get_device(
-        connections={(dr.CONNECTION_NETWORK_MAC, "1234")}
+    device = device_registry.async_get_device_by_connection(
+        (dr.CONNECTION_NETWORK_MAC, "1234"), config_entry.entry_id
     )
     assert device is not None
     assert device.name == expected_device_name