Adjust entity registry access in helper tests (#88965)

This commit is contained in:
epenet 2023-03-01 16:24:31 +01:00 committed by GitHub
parent d65dff3f9e
commit b84eead3f8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 66 additions and 57 deletions

View File

@ -7,10 +7,10 @@ from homeassistant.components.switch import SwitchDeviceClass
from homeassistant.const import ATTR_FRIENDLY_NAME from homeassistant.const import ATTR_FRIENDLY_NAME
from homeassistant.core import Context, HomeAssistant, State from homeassistant.core import Context, HomeAssistant, State
from homeassistant.helpers import ( from homeassistant.helpers import (
area_registry, area_registry as ar,
config_validation as cv, config_validation as cv,
device_registry, device_registry as dr,
entity_registry, entity_registry as er,
intent, intent,
) )
from homeassistant.setup import async_setup_component from homeassistant.setup import async_setup_component
@ -24,12 +24,15 @@ class MockIntentHandler(intent.IntentHandler):
self.slot_schema = slot_schema self.slot_schema = slot_schema
async def test_async_match_states(hass: HomeAssistant) -> None: async def test_async_match_states(
hass: HomeAssistant,
area_registry: ar.AreaRegistry,
entity_registry: er.EntityRegistry,
) -> None:
"""Test async_match_state helper.""" """Test async_match_state helper."""
areas = area_registry.async_get(hass) area_kitchen = area_registry.async_get_or_create("kitchen")
area_kitchen = areas.async_get_or_create("kitchen") area_registry.async_update(area_kitchen.id, aliases={"food room"})
areas.async_update(area_kitchen.id, aliases={"food room"}) area_bedroom = area_registry.async_get_or_create("bedroom")
area_bedroom = areas.async_get_or_create("bedroom")
state1 = State( state1 = State(
"light.kitchen", "on", attributes={ATTR_FRIENDLY_NAME: "kitchen light"} "light.kitchen", "on", attributes={ATTR_FRIENDLY_NAME: "kitchen light"}
@ -39,14 +42,15 @@ async def test_async_match_states(hass: HomeAssistant) -> None:
) )
# Put entities into different areas # Put entities into different areas
entities = entity_registry.async_get(hass) entity_registry.async_get_or_create(
entities.async_get_or_create("light", "demo", "1234", suggested_object_id="kitchen") "light", "demo", "1234", suggested_object_id="kitchen"
entities.async_update_entity(state1.entity_id, area_id=area_kitchen.id) )
entity_registry.async_update_entity(state1.entity_id, area_id=area_kitchen.id)
entities.async_get_or_create( entity_registry.async_get_or_create(
"switch", "demo", "5678", suggested_object_id="bedroom" "switch", "demo", "5678", suggested_object_id="bedroom"
) )
entities.async_update_entity( entity_registry.async_update_entity(
state2.entity_id, state2.entity_id,
area_id=area_bedroom.id, area_id=area_bedroom.id,
device_class=SwitchDeviceClass.OUTLET, device_class=SwitchDeviceClass.OUTLET,
@ -102,17 +106,20 @@ async def test_async_match_states(hass: HomeAssistant) -> None:
) == [state2] ) == [state2]
async def test_match_device_area(hass: HomeAssistant) -> None: async def test_match_device_area(
hass: HomeAssistant,
area_registry: ar.AreaRegistry,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
) -> None:
"""Test async_match_state with a device in an area.""" """Test async_match_state with a device in an area."""
areas = area_registry.async_get(hass) area_kitchen = area_registry.async_get_or_create("kitchen")
area_kitchen = areas.async_get_or_create("kitchen") area_bedroom = area_registry.async_get_or_create("bedroom")
area_bedroom = areas.async_get_or_create("bedroom")
devices = device_registry.async_get(hass) kitchen_device = device_registry.async_get_or_create(
kitchen_device = devices.async_get_or_create(
config_entry_id="1234", connections=set(), identifiers={("demo", "id-1234")} config_entry_id="1234", connections=set(), identifiers={("demo", "id-1234")}
) )
devices.async_update_device(kitchen_device.id, area_id=area_kitchen.id) device_registry.async_update_device(kitchen_device.id, area_id=area_kitchen.id)
state1 = State( state1 = State(
"light.kitchen", "on", attributes={ATTR_FRIENDLY_NAME: "kitchen light"} "light.kitchen", "on", attributes={ATTR_FRIENDLY_NAME: "kitchen light"}
@ -123,12 +130,15 @@ async def test_match_device_area(hass: HomeAssistant) -> None:
state3 = State( state3 = State(
"light.living_room", "on", attributes={ATTR_FRIENDLY_NAME: "living room light"} "light.living_room", "on", attributes={ATTR_FRIENDLY_NAME: "living room light"}
) )
entities = entity_registry.async_get(hass) entity_registry.async_get_or_create(
entities.async_get_or_create("light", "demo", "1234", suggested_object_id="kitchen") "light", "demo", "1234", suggested_object_id="kitchen"
entities.async_update_entity(state1.entity_id, device_id=kitchen_device.id) )
entity_registry.async_update_entity(state1.entity_id, device_id=kitchen_device.id)
entities.async_get_or_create("light", "demo", "5678", suggested_object_id="bedroom") entity_registry.async_get_or_create(
entities.async_update_entity(state2.entity_id, area_id=area_bedroom.id) "light", "demo", "5678", suggested_object_id="bedroom"
)
entity_registry.async_update_entity(state2.entity_id, area_id=area_bedroom.id)
# Match on area/domain # Match on area/domain
assert list( assert list(

View File

@ -21,8 +21,8 @@ from homeassistant.const import (
) )
from homeassistant.core import Context, HomeAssistant, ServiceCall from homeassistant.core import Context, HomeAssistant, ServiceCall
from homeassistant.helpers import ( from homeassistant.helpers import (
device_registry as dev_reg, device_registry as dr,
entity_registry as ent_reg, entity_registry as er,
service, service,
template, template,
) )
@ -96,10 +96,10 @@ def area_mock(hass):
hass.states.async_set("light.Ceiling", STATE_OFF) hass.states.async_set("light.Ceiling", STATE_OFF)
hass.states.async_set("light.Kitchen", STATE_OFF) hass.states.async_set("light.Kitchen", STATE_OFF)
device_in_area = dev_reg.DeviceEntry(area_id="test-area") device_in_area = dr.DeviceEntry(area_id="test-area")
device_no_area = dev_reg.DeviceEntry(id="device-no-area-id") device_no_area = dr.DeviceEntry(id="device-no-area-id")
device_diff_area = dev_reg.DeviceEntry(area_id="diff-area") device_diff_area = dr.DeviceEntry(area_id="diff-area")
device_area_a = dev_reg.DeviceEntry(id="device-area-a-id", area_id="area-a") device_area_a = dr.DeviceEntry(id="device-area-a-id", area_id="area-a")
mock_device_registry( mock_device_registry(
hass, hass,
@ -111,94 +111,94 @@ def area_mock(hass):
}, },
) )
entity_in_own_area = ent_reg.RegistryEntry( entity_in_own_area = er.RegistryEntry(
entity_id="light.in_own_area", entity_id="light.in_own_area",
unique_id="in-own-area-id", unique_id="in-own-area-id",
platform="test", platform="test",
area_id="own-area", area_id="own-area",
) )
config_entity_in_own_area = ent_reg.RegistryEntry( config_entity_in_own_area = er.RegistryEntry(
entity_id="light.config_in_own_area", entity_id="light.config_in_own_area",
unique_id="config-in-own-area-id", unique_id="config-in-own-area-id",
platform="test", platform="test",
area_id="own-area", area_id="own-area",
entity_category=EntityCategory.CONFIG, entity_category=EntityCategory.CONFIG,
) )
hidden_entity_in_own_area = ent_reg.RegistryEntry( hidden_entity_in_own_area = er.RegistryEntry(
entity_id="light.hidden_in_own_area", entity_id="light.hidden_in_own_area",
unique_id="hidden-in-own-area-id", unique_id="hidden-in-own-area-id",
platform="test", platform="test",
area_id="own-area", area_id="own-area",
hidden_by=ent_reg.RegistryEntryHider.USER, hidden_by=er.RegistryEntryHider.USER,
) )
entity_in_area = ent_reg.RegistryEntry( entity_in_area = er.RegistryEntry(
entity_id="light.in_area", entity_id="light.in_area",
unique_id="in-area-id", unique_id="in-area-id",
platform="test", platform="test",
device_id=device_in_area.id, device_id=device_in_area.id,
) )
config_entity_in_area = ent_reg.RegistryEntry( config_entity_in_area = er.RegistryEntry(
entity_id="light.config_in_area", entity_id="light.config_in_area",
unique_id="config-in-area-id", unique_id="config-in-area-id",
platform="test", platform="test",
device_id=device_in_area.id, device_id=device_in_area.id,
entity_category=EntityCategory.CONFIG, entity_category=EntityCategory.CONFIG,
) )
hidden_entity_in_area = ent_reg.RegistryEntry( hidden_entity_in_area = er.RegistryEntry(
entity_id="light.hidden_in_area", entity_id="light.hidden_in_area",
unique_id="hidden-in-area-id", unique_id="hidden-in-area-id",
platform="test", platform="test",
device_id=device_in_area.id, device_id=device_in_area.id,
hidden_by=ent_reg.RegistryEntryHider.USER, hidden_by=er.RegistryEntryHider.USER,
) )
entity_in_other_area = ent_reg.RegistryEntry( entity_in_other_area = er.RegistryEntry(
entity_id="light.in_other_area", entity_id="light.in_other_area",
unique_id="in-area-a-id", unique_id="in-area-a-id",
platform="test", platform="test",
device_id=device_in_area.id, device_id=device_in_area.id,
area_id="other-area", area_id="other-area",
) )
entity_assigned_to_area = ent_reg.RegistryEntry( entity_assigned_to_area = er.RegistryEntry(
entity_id="light.assigned_to_area", entity_id="light.assigned_to_area",
unique_id="assigned-area-id", unique_id="assigned-area-id",
platform="test", platform="test",
device_id=device_in_area.id, device_id=device_in_area.id,
area_id="test-area", area_id="test-area",
) )
entity_no_area = ent_reg.RegistryEntry( entity_no_area = er.RegistryEntry(
entity_id="light.no_area", entity_id="light.no_area",
unique_id="no-area-id", unique_id="no-area-id",
platform="test", platform="test",
device_id=device_no_area.id, device_id=device_no_area.id,
) )
config_entity_no_area = ent_reg.RegistryEntry( config_entity_no_area = er.RegistryEntry(
entity_id="light.config_no_area", entity_id="light.config_no_area",
unique_id="config-no-area-id", unique_id="config-no-area-id",
platform="test", platform="test",
device_id=device_no_area.id, device_id=device_no_area.id,
entity_category=EntityCategory.CONFIG, entity_category=EntityCategory.CONFIG,
) )
hidden_entity_no_area = ent_reg.RegistryEntry( hidden_entity_no_area = er.RegistryEntry(
entity_id="light.hidden_no_area", entity_id="light.hidden_no_area",
unique_id="hidden-no-area-id", unique_id="hidden-no-area-id",
platform="test", platform="test",
device_id=device_no_area.id, device_id=device_no_area.id,
hidden_by=ent_reg.RegistryEntryHider.USER, hidden_by=er.RegistryEntryHider.USER,
) )
entity_diff_area = ent_reg.RegistryEntry( entity_diff_area = er.RegistryEntry(
entity_id="light.diff_area", entity_id="light.diff_area",
unique_id="diff-area-id", unique_id="diff-area-id",
platform="test", platform="test",
device_id=device_diff_area.id, device_id=device_diff_area.id,
) )
entity_in_area_a = ent_reg.RegistryEntry( entity_in_area_a = er.RegistryEntry(
entity_id="light.in_area_a", entity_id="light.in_area_a",
unique_id="in-area-a-id", unique_id="in-area-a-id",
platform="test", platform="test",
device_id=device_area_a.id, device_id=device_area_a.id,
area_id="area-a", area_id="area-a",
) )
entity_in_area_b = ent_reg.RegistryEntry( entity_in_area_b = er.RegistryEntry(
entity_id="light.in_area_b", entity_id="light.in_area_b",
unique_id="in-area-b-id", unique_id="in-area-b-id",
platform="test", platform="test",
@ -403,11 +403,12 @@ class TestServiceHelpers(unittest.TestCase):
assert mock_log.call_count == 3 assert mock_log.call_count == 3
async def test_service_call_entry_id(hass: HomeAssistant) -> None: async def test_service_call_entry_id(
hass: HomeAssistant, entity_registry: er.EntityRegistry
) -> None:
"""Test service call with entity specified by entity registry ID.""" """Test service call with entity specified by entity registry ID."""
registry = ent_reg.async_get(hass)
calls = async_mock_service(hass, "test_domain", "test_service") calls = async_mock_service(hass, "test_domain", "test_service")
entry = registry.async_get_or_create( entry = entity_registry.async_get_or_create(
"hello", "hue", "1234", suggested_object_id="world" "hello", "hue", "1234", suggested_object_id="world"
) )
@ -946,7 +947,7 @@ async def test_domain_control_unauthorized(
mock_registry( mock_registry(
hass, hass,
{ {
"light.kitchen": ent_reg.RegistryEntry( "light.kitchen": er.RegistryEntry(
entity_id="light.kitchen", entity_id="light.kitchen",
unique_id="kitchen", unique_id="kitchen",
platform="test_domain", platform="test_domain",
@ -987,7 +988,7 @@ async def test_domain_control_admin(
mock_registry( mock_registry(
hass, hass,
{ {
"light.kitchen": ent_reg.RegistryEntry( "light.kitchen": er.RegistryEntry(
entity_id="light.kitchen", entity_id="light.kitchen",
unique_id="kitchen", unique_id="kitchen",
platform="test_domain", platform="test_domain",
@ -1025,7 +1026,7 @@ async def test_domain_control_no_user(hass: HomeAssistant) -> None:
mock_registry( mock_registry(
hass, hass,
{ {
"light.kitchen": ent_reg.RegistryEntry( "light.kitchen": er.RegistryEntry(
entity_id="light.kitchen", entity_id="light.kitchen",
unique_id="kitchen", unique_id="kitchen",
platform="test_domain", platform="test_domain",
@ -1213,9 +1214,7 @@ async def test_async_extract_entities_warn_referenced(
async def test_async_extract_config_entry_ids(hass: HomeAssistant) -> None: async def test_async_extract_config_entry_ids(hass: HomeAssistant) -> None:
"""Test we can find devices that have no entities.""" """Test we can find devices that have no entities."""
device_no_entities = dev_reg.DeviceEntry( device_no_entities = dr.DeviceEntry(id="device-no-entities", config_entries={"abc"})
id="device-no-entities", config_entries={"abc"}
)
call = ServiceCall( call = ServiceCall(
"homeassistant", "homeassistant",