evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add support to entity registry for overriding device_class (#59985)

Browse source
This commit is contained in:
Erik Montnemery 2021-11-22 17:38:06 +01:00 • committed by GitHub
parent 67e13b35db
commit 49a27e12ad
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
31 changed files with 136 additions and 69 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
homeassistant/components/config/entity_registry.py
View file

@ -70,6 +70,7 @@ async def websocket_get_entity(hass, connection, msg):
vol.Required("entity_id"): cv.entity_id,
# If passed in, we update value. Passing None will remove old value.
vol.Optional("area_id"): vol.Any(str, None),
vol.Optional("device_class"): vol.Any(str, None),
vol.Optional("icon"): vol.Any(str, None),
vol.Optional("name"): vol.Any(str, None),
vol.Optional("new_entity_id"): str,
@ -92,7 +93,7 @@ async def websocket_update_entity(hass, connection, msg):
changes = {}
for key in ("area_id", "disabled_by", "icon", "name"):
for key in ("area_id", "device_class", "disabled_by", "icon", "name"):
if key in msg:
changes[key] = msg[key]
@ -185,6 +186,8 @@ def _entry_ext_dict(entry):
"""Convert entry to API format."""
data = _entry_dict(entry)
data["capabilities"] = entry.capabilities
data["device_class"] = entry.device_class
data["original_device_class"] = entry.original_device_class
data["original_icon"] = entry.original_icon
data["original_name"] = entry.original_name
data["unique_id"] = entry.unique_id

2
homeassistant/components/homekit/__init__.py
View file

@ -859,7 +859,7 @@ class HomeKit:
ent_reg_ent is None
or ent_reg_ent.device_id is None
or ent_reg_ent.device_id not in device_lookup
or ent_reg_ent.device_class
or (ent_reg_ent.device_class or ent_reg_ent.original_device_class)
in (DEVICE_CLASS_BATTERY_CHARGING, DEVICE_CLASS_BATTERY)
):
return

2
homeassistant/components/hue/migration.py
View file

@ -126,7 +126,7 @@ async def handle_v2_migration(hass: core.HomeAssistant, entry: ConfigEntry) -> N
continue
# migrate sensors
matched_dev_class = sensor_class_mapping.get(
ent.device_class or "unknown"
ent.original_device_class or "unknown"
)
if matched_dev_class is None:
# this may happen if we're looking at orphaned or unsupported entity

2
homeassistant/components/mobile_app/binary_sensor.py
View file

@ -35,7 +35,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
continue
config = {
ATTR_SENSOR_ATTRIBUTES: {},
ATTR_SENSOR_DEVICE_CLASS: entry.device_class,
ATTR_SENSOR_DEVICE_CLASS: entry.device_class or entry.original_device_class,
ATTR_SENSOR_ICON: entry.original_icon,
ATTR_SENSOR_NAME: entry.original_name,
ATTR_SENSOR_STATE: None,

2
homeassistant/components/mobile_app/sensor.py
View file

@ -46,7 +46,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
continue
config = {
ATTR_SENSOR_ATTRIBUTES: {},
ATTR_SENSOR_DEVICE_CLASS: entry.device_class,
ATTR_SENSOR_DEVICE_CLASS: entry.device_class or entry.original_device_class,
ATTR_SENSOR_ICON: entry.original_icon,
ATTR_SENSOR_NAME: entry.original_name,
ATTR_SENSOR_STATE: None,

2
homeassistant/components/shelly/entity.py
View file

@ -135,7 +135,7 @@ async def async_restore_block_attribute_entities(
name="",
icon=entry.original_icon,
unit=entry.unit_of_measurement,
device_class=entry.device_class,
device_class=entry.original_device_class,
)
entities.append(

6
homeassistant/helpers/entity.py
View file

@ -127,7 +127,7 @@ def get_device_class(hass: HomeAssistant, entity_id: str) -> str | None:
if not (entry := entity_registry.async_get(entity_id)):
raise HomeAssistantError(f"Unknown entity {entity_id}")
return entry.device_class
return entry.device_class or entry.original_device_class
def get_supported_features(hass: HomeAssistant, entity_id: str) -> int:
@ -540,7 +540,9 @@ class Entity(ABC):
if (attribution := self.attribution) is not None:
attr[ATTR_ATTRIBUTION] = attribution
if (device_class := self.device_class) is not None:
if (
device_class := (entry and entry.device_class) or self.device_class
) is not None:
attr[ATTR_DEVICE_CLASS] = str(device_class)
if (entity_picture := self.entity_picture) is not None:

2
homeassistant/helpers/entity_platform.py
View file

@ -511,11 +511,11 @@ class EntityPlatform:
entity.unique_id,
capabilities=entity.capability_attributes,
config_entry=self.config_entry,
device_class=entity.device_class,
device_id=device_id,
disabled_by=disabled_by,
entity_category=entity.entity_category,
known_object_ids=self.entities.keys(),
original_device_class=entity.device_class,
original_icon=entity.icon,
original_name=entity.name,
suggested_object_id=suggested_object_id,

25
homeassistant/helpers/entity_registry.py
View file

@ -105,6 +105,7 @@ class RegistryEntry:
icon: str | None = attr.ib(default=None)
name: str | None = attr.ib(default=None)
# As set by integration
original_device_class: str | None = attr.ib(default=None)
original_icon: str | None = attr.ib(default=None)
original_name: str | None = attr.ib(default=None)
supported_features: int = attr.ib(default=0)
@ -128,8 +129,9 @@ class RegistryEntry:
if self.capabilities is not None:
attrs.update(self.capabilities)
if self.device_class is not None:
attrs[ATTR_DEVICE_CLASS] = self.device_class
device_class = self.device_class or self.original_device_class
if device_class is not None:
attrs[ATTR_DEVICE_CLASS] = device_class
icon = self.icon or self.original_icon
if icon is not None:
@ -190,7 +192,8 @@ class EntityRegistry:
for entity in self.entities.values():
if not entity.device_id:
continue
domain_device_class = (entity.domain, entity.device_class)
device_class = entity.device_class or entity.original_device_class
domain_device_class = (entity.domain, device_class)
if domain_device_class not in domain_device_classes:
continue
if entity.device_id not in lookup:
@ -268,9 +271,9 @@ class EntityRegistry:
area_id: str | None = None,
capabilities: Mapping[str, Any] | None = None,
config_entry: ConfigEntry | None = None,
device_class: str | None = None,
device_id: str | None = None,
entity_category: str | None = None,
original_device_class: str | None = None,
original_icon: str | None = None,
original_name: str | None = None,
supported_features: int | None = None,
@ -289,9 +292,9 @@ class EntityRegistry:
area_id=area_id or UNDEFINED,
capabilities=capabilities or UNDEFINED,
config_entry_id=config_entry_id or UNDEFINED,
device_class=device_class or UNDEFINED,
device_id=device_id or UNDEFINED,
entity_category=entity_category or UNDEFINED,
original_device_class=original_device_class or UNDEFINED,
original_icon=original_icon or UNDEFINED,
original_name=original_name or UNDEFINED,
supported_features=supported_features or UNDEFINED,
@ -320,11 +323,11 @@ class EntityRegistry:
area_id=area_id,
capabilities=capabilities,
config_entry_id=config_entry_id,
device_class=device_class,
device_id=device_id,
disabled_by=disabled_by,
entity_category=entity_category,
entity_id=entity_id,
original_device_class=original_device_class,
original_icon=original_icon,
original_name=original_name,
platform=platform,
@ -409,6 +412,7 @@ class EntityRegistry:
name: str | None | UndefinedType = UNDEFINED,
new_entity_id: str | UndefinedType = UNDEFINED,
new_unique_id: str | UndefinedType = UNDEFINED,
original_device_class: str | None | UndefinedType = UNDEFINED,
original_icon: str | None | UndefinedType = UNDEFINED,
original_name: str | None | UndefinedType = UNDEFINED,
unit_of_measurement: str | None | UndefinedType = UNDEFINED,
@ -425,6 +429,7 @@ class EntityRegistry:
name=name,
new_entity_id=new_entity_id,
new_unique_id=new_unique_id,
original_device_class=original_device_class,
original_icon=original_icon,
original_name=original_name,
unit_of_measurement=unit_of_measurement,
@ -446,6 +451,7 @@ class EntityRegistry:
name: str | None | UndefinedType = UNDEFINED,
new_entity_id: str | UndefinedType = UNDEFINED,
new_unique_id: str | UndefinedType = UNDEFINED,
original_device_class: str | None | UndefinedType = UNDEFINED,
original_icon: str | None | UndefinedType = UNDEFINED,
original_name: str | None | UndefinedType = UNDEFINED,
supported_features: int | UndefinedType = UNDEFINED,
@ -467,6 +473,7 @@ class EntityRegistry:
("entity_category", entity_category),
("icon", icon),
("name", name),
("original_device_class", original_device_class),
("original_icon", original_icon),
("original_name", original_name),
("supported_features", supported_features),
@ -552,6 +559,7 @@ class EntityRegistry:
entity_id=entity["entity_id"],
icon=entity["icon"],
name=entity["name"],
original_device_class=entity["original_device_class"],
original_icon=entity["original_icon"],
original_name=entity["original_name"],
platform=entity["platform"],
@ -585,6 +593,7 @@ class EntityRegistry:
"entity_id": entry.entity_id,
"icon": entry.icon,
"name": entry.name,
"original_device_class": entry.original_device_class,
"original_icon": entry.original_icon,
"original_name": entry.original_name,
"platform": entry.platform,
@ -741,6 +750,10 @@ async def _async_migrate(
entity["platform"] = entity["platform"]
entity["supported_features"] = entity.get("supported_features", 0)
entity["unit_of_measurement"] = entity.get("unit_of_measurement")
# Move device_class to original_device_class
entity["original_device_class"] = entity["device_class"]
entity["device_class"] = None
if old_major_version > 1:
raise NotImplementedError
return data

2
tests/components/binary_sensor/test_device_condition.py
View file

@ -93,7 +93,7 @@ async def test_get_conditions_no_state(hass, device_reg, entity_reg):
"test",
f"5678_{device_class}",
device_id=device_entry.id,
device_class=device_class,
original_device_class=device_class,
).entity_id
await hass.async_block_till_done()

2
tests/components/binary_sensor/test_device_trigger.py
View file

@ -96,7 +96,7 @@ async def test_get_triggers_no_state(hass, device_reg, entity_reg):
"test",
f"5678_{device_class}",
device_id=device_entry.id,
device_class=device_class,
original_device_class=device_class,
).entity_id
await hass.async_block_till_done()

2
tests/components/blebox/test_cover.py
View file

@ -136,7 +136,7 @@ async def test_init_shutterbox(shutterbox, hass, config):
state = hass.states.get(entity_id)
assert state.name == "shutterBox-position"
assert entry.device_class == DEVICE_CLASS_SHUTTER
assert entry.original_device_class == DEVICE_CLASS_SHUTTER
supported_features = state.attributes[ATTR_SUPPORTED_FEATURES]
assert supported_features & SUPPORT_OPEN

4
tests/components/canary/test_sensor.py
View file

@ -78,7 +78,7 @@ async def test_sensors_pro(hass, canary) -> None:
for (sensor_id, data) in sensors.items():
entity_entry = registry.async_get(f"sensor.{sensor_id}")
assert entity_entry
assert entity_entry.device_class == data[3]
assert entity_entry.original_device_class == data[3]
assert entity_entry.unique_id == data[0]
assert entity_entry.original_icon == data[4]
@ -197,7 +197,7 @@ async def test_sensors_flex(hass, canary) -> None:
for (sensor_id, data) in sensors.items():
entity_entry = registry.async_get(f"sensor.{sensor_id}")
assert entity_entry
assert entity_entry.device_class == data[3]
assert entity_entry.original_device_class == data[3]
assert entity_entry.unique_id == data[0]
assert entity_entry.original_icon == data[4]

21
tests/components/config/test_entity_registry.py
View file

@ -103,12 +103,14 @@ async def test_get_entity(hass, client):
"area_id": None,
"capabilities": None,
"config_entry_id": None,
"device_class": None,
"device_id": None,
"disabled_by": None,
"entity_category": None,
"entity_id": "test_domain.name",
"icon": None,
"name": "Hello World",
"original_device_class": None,
"original_icon": None,
"original_name": None,
"platform": "test_platform",
@ -128,12 +130,14 @@ async def test_get_entity(hass, client):
"area_id": None,
"capabilities": None,
"config_entry_id": None,
"device_class": None,
"device_id": None,
"disabled_by": None,
"entity_category": None,
"entity_id": "test_domain.no_name",
"icon": None,
"name": None,
"original_device_class": None,
"original_icon": None,
"original_name": None,
"platform": "test_platform",
@ -165,15 +169,16 @@ async def test_update_entity(hass, client):
assert state.name == "before update"
assert state.attributes[ATTR_ICON] == "icon:before update"
# UPDATE NAME & ICON & AREA
# UPDATE AREA, DEVICE_CLASS, ICON AND NAME
await client.send_json(
{
"id": 6,
"type": "config/entity_registry/update",
"entity_id": "test_domain.world",
"name": "after update",
"icon": "icon:after update",
"area_id": "mock-area-id",
"device_class": "custom_device_class",
"icon": "icon:after update",
"name": "after update",
}
)
@ -184,12 +189,14 @@ async def test_update_entity(hass, client):
"area_id": "mock-area-id",
"capabilities": None,
"config_entry_id": None,
"device_class": "custom_device_class",
"device_id": None,
"disabled_by": None,
"entity_category": None,
"entity_id": "test_domain.world",
"icon": "icon:after update",
"name": "after update",
"original_device_class": None,
"original_icon": None,
"original_name": None,
"platform": "test_platform",
@ -233,12 +240,14 @@ async def test_update_entity(hass, client):
"area_id": "mock-area-id",
"capabilities": None,
"config_entry_id": None,
"device_class": "custom_device_class",
"device_id": None,
"disabled_by": None,
"entity_category": None,
"entity_id": "test_domain.world",
"icon": "icon:after update",
"name": "after update",
"original_device_class": None,
"original_icon": None,
"original_name": None,
"platform": "test_platform",
@ -288,12 +297,14 @@ async def test_update_entity_require_restart(hass, client):
"area_id": None,
"capabilities": None,
"config_entry_id": config_entry.entry_id,
"device_class": None,
"device_id": None,
"disabled_by": None,
"entity_category": None,
"entity_id": "test_domain.world",
"icon": None,
"name": None,
"original_device_class": None,
"original_icon": None,
"original_name": None,
"platform": "test_platform",
@ -390,12 +401,14 @@ async def test_update_entity_no_changes(hass, client):
"area_id": None,
"capabilities": None,
"config_entry_id": None,
"device_class": None,
"device_id": None,
"disabled_by": None,
"entity_category": None,
"entity_id": "test_domain.world",
"icon": None,
"name": "name of entity",
"original_device_class": None,
"original_icon": None,
"original_name": None,
"platform": "test_platform",
@ -471,12 +484,14 @@ async def test_update_entity_id(hass, client):
"area_id": None,
"capabilities": None,
"config_entry_id": None,
"device_class": None,
"device_id": None,
"disabled_by": None,
"entity_category": None,
"entity_id": "test_domain.planet",
"icon": None,
"name": None,
"original_device_class": None,
"original_icon": None,
"original_name": None,
"platform": "test_platform",

6
tests/components/directv/test_media_player.py
View file

@ -160,15 +160,15 @@ async def test_unique_id(
entity_registry = er.async_get(hass)
main = entity_registry.async_get(MAIN_ENTITY_ID)
assert main.device_class == DEVICE_CLASS_RECEIVER
assert main.original_device_class == DEVICE_CLASS_RECEIVER
assert main.unique_id == "028877455858"
client = entity_registry.async_get(CLIENT_ENTITY_ID)
assert client.device_class == DEVICE_CLASS_RECEIVER
assert client.original_device_class == DEVICE_CLASS_RECEIVER
assert client.unique_id == "2CA17D1CD30X"
unavailable_client = entity_registry.async_get(UNAVAILABLE_ENTITY_ID)
assert unavailable_client.device_class == DEVICE_CLASS_RECEIVER
assert unavailable_client.original_device_class == DEVICE_CLASS_RECEIVER
assert unavailable_client.unique_id == "9XXXXXXXXXX9"

6
tests/components/greeneye_monitor/conftest.py
View file

@ -45,7 +45,7 @@ def assert_temperature_sensor_registered(
):
"""Assert that a temperature sensor entity was registered properly."""
sensor = assert_sensor_registered(hass, serial_number, "temp", number, name)
assert sensor.device_class == DEVICE_CLASS_TEMPERATURE
assert sensor.original_device_class == DEVICE_CLASS_TEMPERATURE
def assert_pulse_counter_registered(
@ -67,7 +67,7 @@ def assert_power_sensor_registered(
"""Assert that a power sensor entity was registered properly."""
sensor = assert_sensor_registered(hass, serial_number, "current", number, name)
assert sensor.unit_of_measurement == POWER_WATT
assert sensor.device_class == DEVICE_CLASS_POWER
assert sensor.original_device_class == DEVICE_CLASS_POWER
def assert_voltage_sensor_registered(
@ -76,7 +76,7 @@ def assert_voltage_sensor_registered(
"""Assert that a voltage sensor entity was registered properly."""
sensor = assert_sensor_registered(hass, serial_number, "volts", number, name)
assert sensor.unit_of_measurement == ELECTRIC_POTENTIAL_VOLT
assert sensor.device_class == DEVICE_CLASS_VOLTAGE
assert sensor.original_device_class == DEVICE_CLASS_VOLTAGE
def assert_sensor_registered(

2
tests/components/group/test_binary_sensor.py
View file

@ -40,7 +40,7 @@ async def test_default_state(hass):
assert entry
assert entry.unique_id == "unique_identifier"
assert entry.original_name == "Bedroom Group"
assert entry.device_class == "presence"
assert entry.original_device_class == "presence"
async def test_state_reporting_all(hass):

16
tests/components/homekit/test_homekit.py
View file

@ -1118,14 +1118,14 @@ async def test_homekit_finds_linked_batteries(
"powerwall",
"battery_charging",
device_id=device_entry.id,
device_class=DEVICE_CLASS_BATTERY_CHARGING,
original_device_class=DEVICE_CLASS_BATTERY_CHARGING,
)
battery_sensor = entity_reg.async_get_or_create(
"sensor",
"powerwall",
"battery",
device_id=device_entry.id,
device_class=DEVICE_CLASS_BATTERY,
original_device_class=DEVICE_CLASS_BATTERY,
)
light = entity_reg.async_get_or_create(
"light", "powerwall", "demo", device_id=device_entry.id
@ -1187,14 +1187,14 @@ async def test_homekit_async_get_integration_fails(
"invalid_integration_does_not_exist",
"battery_charging",
device_id=device_entry.id,
device_class=DEVICE_CLASS_BATTERY_CHARGING,
original_device_class=DEVICE_CLASS_BATTERY_CHARGING,
)
battery_sensor = entity_reg.async_get_or_create(
"sensor",
"invalid_integration_does_not_exist",
"battery",
device_id=device_entry.id,
device_class=DEVICE_CLASS_BATTERY,
original_device_class=DEVICE_CLASS_BATTERY,
)
light = entity_reg.async_get_or_create(
"light", "invalid_integration_does_not_exist", "demo", device_id=device_entry.id
@ -1334,14 +1334,14 @@ async def test_homekit_ignored_missing_devices(
"powerwall",
"battery_charging",
device_id=device_entry.id,
device_class=DEVICE_CLASS_BATTERY_CHARGING,
original_device_class=DEVICE_CLASS_BATTERY_CHARGING,
)
entity_reg.async_get_or_create(
"sensor",
"powerwall",
"battery",
device_id=device_entry.id,
device_class=DEVICE_CLASS_BATTERY,
original_device_class=DEVICE_CLASS_BATTERY,
)
light = entity_reg.async_get_or_create(
"light", "powerwall", "demo", device_id=device_entry.id
@ -1404,7 +1404,7 @@ async def test_homekit_finds_linked_motion_sensors(
"camera",
"motion_sensor",
device_id=device_entry.id,
device_class=DEVICE_CLASS_MOTION,
original_device_class=DEVICE_CLASS_MOTION,
)
camera = entity_reg.async_get_or_create(
"camera", "camera", "demo", device_id=device_entry.id
@ -1466,7 +1466,7 @@ async def test_homekit_finds_linked_humidity_sensors(
"humidifier",
"humidity_sensor",
device_id=device_entry.id,
device_class=DEVICE_CLASS_HUMIDITY,
original_device_class=DEVICE_CLASS_HUMIDITY,
)
humidifier = entity_reg.async_get_or_create(
"humidifier", "humidifier", "demo", device_id=device_entry.id

4
tests/components/homekit/test_type_covers.py
View file

@ -506,7 +506,7 @@ async def test_windowcovering_basic_restore(hass, hk_driver, events):
suggested_object_id="all_info_set",
capabilities={},
supported_features=SUPPORT_STOP,
device_class="mock-device-class",
original_device_class="mock-device-class",
)
hass.bus.async_fire(EVENT_HOMEASSISTANT_START, {})
@ -544,7 +544,7 @@ async def test_windowcovering_restore(hass, hk_driver, events):
suggested_object_id="all_info_set",
capabilities={},
supported_features=SUPPORT_STOP,
device_class="mock-device-class",
original_device_class="mock-device-class",
)
hass.bus.async_fire(EVENT_HOMEASSISTANT_START, {})

2
tests/components/homekit/test_type_fans.py
View file

@ -546,7 +546,7 @@ async def test_fan_restore(hass, hk_driver, events):
suggested_object_id="all_info_set",
capabilities={"speed_list": ["off", "low", "medium", "high"]},
supported_features=SUPPORT_SET_SPEED | SUPPORT_OSCILLATE | SUPPORT_DIRECTION,
device_class="mock-device-class",
original_device_class="mock-device-class",
)
hass.bus.async_fire(EVENT_HOMEASSISTANT_START, {})

2
tests/components/homekit/test_type_lights.py
View file

@ -546,7 +546,7 @@ async def test_light_restore(hass, hk_driver, events):
suggested_object_id="all_info_set",
capabilities={"supported_color_modes": ["brightness"], "max": 100},
supported_features=5,
device_class="mock-device-class",
original_device_class="mock-device-class",
)
hass.bus.async_fire(EVENT_HOMEASSISTANT_START, {})

4
tests/components/homekit/test_type_media_players.py
View file

@ -431,7 +431,7 @@ async def test_tv_restore(hass, hk_driver, events):
"generic",
"1234",
suggested_object_id="simple",
device_class=DEVICE_CLASS_TV,
original_device_class=DEVICE_CLASS_TV,
)
registry.async_get_or_create(
"media_player",
@ -442,7 +442,7 @@ async def test_tv_restore(hass, hk_driver, events):
ATTR_INPUT_SOURCE_LIST: ["HDMI 1", "HDMI 2", "HDMI 3", "HDMI 4"],
},
supported_features=3469,
device_class=DEVICE_CLASS_TV,
original_device_class=DEVICE_CLASS_TV,
)
hass.bus.async_fire(EVENT_HOMEASSISTANT_START, {})

4
tests/components/homekit/test_type_sensors.py
View file

@ -333,14 +333,14 @@ async def test_sensor_restore(hass, hk_driver, events):
"generic",
"1234",
suggested_object_id="temperature",
device_class="temperature",
original_device_class="temperature",
)
registry.async_get_or_create(
"sensor",
"generic",
"12345",
suggested_object_id="humidity",
device_class="humidity",
original_device_class="humidity",
unit_of_measurement=PERCENTAGE,
)
hass.bus.async_fire(EVENT_HOMEASSISTANT_START, {})

4
tests/components/homekit/test_type_thermostats.py
View file

@ -1018,7 +1018,7 @@ async def test_thermostat_restore(hass, hk_driver, events):
ATTR_HVAC_MODES: [HVAC_MODE_HEAT_COOL, HVAC_MODE_OFF],
},
supported_features=0,
device_class="mock-device-class",
original_device_class="mock-device-class",
)
hass.bus.async_fire(EVENT_HOMEASSISTANT_START, {})
@ -1826,7 +1826,7 @@ async def test_water_heater_restore(hass, hk_driver, events):
suggested_object_id="all_info_set",
capabilities={ATTR_MIN_TEMP: 60, ATTR_MAX_TEMP: 70},
supported_features=0,
device_class="mock-device-class",
original_device_class="mock-device-class",
)
hass.bus.async_fire(EVENT_HOMEASSISTANT_START, {})

2
tests/components/hue/test_migration.py
View file

@ -118,7 +118,7 @@ async def test_sensor_entity_migration(
f"{device_mac}-{dev_class}",
suggested_object_id=f"hue_migrated_{dev_class}_sensor",
device_id=device.id,
device_class=dev_class,
original_device_class=dev_class,
)
# now run the migration and check results

2
tests/components/nzbget/test_sensor.py
View file

@ -45,7 +45,7 @@ async def test_sensors(hass, nzbget_api) -> None:
for (sensor_id, data) in sensors.items():
entity_entry = registry.async_get(f"sensor.nzbgettest_{sensor_id}")
assert entity_entry
assert entity_entry.device_class == data[3]
assert entity_entry.original_device_class == data[3]
assert entity_entry.unique_id == f"{entry.entry_id}_{data[0]}"
state = hass.states.get(f"sensor.nzbgettest_{sensor_id}")

4
tests/components/roku/test_media_player.py
View file

@ -89,7 +89,7 @@ async def test_setup(hass: HomeAssistant, aioclient_mock: AiohttpClientMocker) -
assert hass.states.get(MAIN_ENTITY_ID)
assert main
assert main.device_class == DEVICE_CLASS_RECEIVER
assert main.original_device_class == DEVICE_CLASS_RECEIVER
assert main.unique_id == UPNP_SERIAL
@ -121,7 +121,7 @@ async def test_tv_setup(
assert hass.states.get(TV_ENTITY_ID)
assert tv
assert tv.device_class == DEVICE_CLASS_TV
assert tv.original_device_class == DEVICE_CLASS_TV
assert tv.unique_id == TV_SERIAL

4
tests/components/sensor/test_device_condition.py
View file

@ -100,7 +100,7 @@ async def test_get_conditions_no_state(hass, device_reg, entity_reg):
"test",
f"5678_{device_class}",
device_id=device_entry.id,
device_class=device_class,
original_device_class=device_class,
unit_of_measurement=UNITS_OF_MEASUREMENT.get(device_class),
).entity_id
@ -155,7 +155,7 @@ async def test_get_condition_capabilities(
"test",
platform.ENTITIES["battery"].unique_id,
device_id=device_entry.id,
device_class=device_class_reg,
original_device_class=device_class_reg,
unit_of_measurement=unit_reg,
).entity_id
if set_state:

2
tests/components/sensor/test_device_trigger.py
View file

@ -122,7 +122,7 @@ async def test_get_trigger_capabilities(
"test",
platform.ENTITIES["battery"].unique_id,
device_id=device_entry.id,
device_class=device_class_reg,
original_device_class=device_class_reg,
unit_of_measurement=unit_reg,
).entity_id
if set_state:

3
tests/helpers/test_entity_platform.py
View file

@ -1105,10 +1105,11 @@ async def test_entity_info_added_to_entity_registry(hass):
"default",
"test_domain",
capabilities={"max": 100},
device_class="mock-device-class",
device_class=None,
entity_category="config",
icon=None,
name=None,
original_device_class="mock-device-class",
original_icon="nice:icon",
original_name="best name",
supported_features=5,

59
tests/helpers/test_entity_registry.py
View file

@ -74,10 +74,10 @@ def test_get_or_create_updates_data(registry):
area_id="mock-area-id",
capabilities={"max": 100},
config_entry=orig_config_entry,
device_class="mock-device-class",
device_id="mock-dev-id",
disabled_by=er.DISABLED_HASS,
entity_category="config",
original_device_class="mock-device-class",
original_icon="initial-original_icon",
original_name="initial-original_name",
supported_features=5,
@ -91,12 +91,13 @@ def test_get_or_create_updates_data(registry):
area_id="mock-area-id",
capabilities={"max": 100},
config_entry_id=orig_config_entry.entry_id,
device_class="mock-device-class",
device_class=None,
device_id="mock-dev-id",
disabled_by=er.DISABLED_HASS,
entity_category="config",
icon=None,
name=None,
original_device_class="mock-device-class",
original_icon="initial-original_icon",
original_name="initial-original_name",
supported_features=5,
@ -112,10 +113,10 @@ def test_get_or_create_updates_data(registry):
area_id="new-mock-area-id",
capabilities={"new-max": 100},
config_entry=new_config_entry,
device_class="new-mock-device-class",
device_id="new-mock-dev-id",
disabled_by=er.DISABLED_USER,
entity_category=None,
original_device_class="new-mock-device-class",
original_icon="updated-original_icon",
original_name="updated-original_name",
supported_features=10,
@ -129,12 +130,13 @@ def test_get_or_create_updates_data(registry):
area_id="new-mock-area-id",
capabilities={"new-max": 100},
config_entry_id=new_config_entry.entry_id,
device_class="new-mock-device-class",
device_class=None,
device_id="new-mock-dev-id",
disabled_by=er.DISABLED_HASS, # Should not be updated
entity_category="config",
icon=None,
name=None,
original_device_class="new-mock-device-class",
original_icon="updated-original_icon",
original_name="updated-original_name",
supported_features=10,
@ -182,17 +184,20 @@ async def test_loading_saving_data(hass, registry):
area_id="mock-area-id",
capabilities={"max": 100},
config_entry=mock_config,
device_class="mock-device-class",
device_id="mock-dev-id",
disabled_by=er.DISABLED_HASS,
entity_category="config",
original_device_class="mock-device-class",
original_icon="hass:original-icon",
original_name="Original Name",
supported_features=5,
unit_of_measurement="initial-unit_of_measurement",
)
orig_entry2 = registry.async_update_entity(
orig_entry2.entity_id, name="User Name", icon="hass:user-icon"
orig_entry2.entity_id,
device_class="user-class",
name="User Name",
icon="hass:user-icon",
)
assert len(registry.entities) == 2
@ -213,12 +218,13 @@ async def test_loading_saving_data(hass, registry):
assert new_entry2.area_id == "mock-area-id"
assert new_entry2.capabilities == {"max": 100}
assert new_entry2.config_entry_id == mock_config.entry_id
assert new_entry2.device_class == "mock-device-class"
assert new_entry2.device_class == "user-class"
assert new_entry2.device_id == "mock-dev-id"
assert new_entry2.disabled_by == er.DISABLED_HASS
assert new_entry2.entity_category == "config"
assert new_entry2.icon == "hass:user-icon"
assert new_entry2.name == "User Name"
assert new_entry2.original_device_class == "mock-device-class"
assert new_entry2.original_icon == "hass:original-icon"
assert new_entry2.original_name == "Original Name"
assert new_entry2.supported_features == 5
@ -411,6 +417,33 @@ async def test_migration_yaml_to_json(hass):
assert entry.config_entry_id == "test-config-id"
@pytest.mark.parametrize("load_registries", [False])
async def test_migration_1_1_to_1_2(hass, hass_storage):
"""Test migration from version 1.1 to 1.2."""
hass_storage[er.STORAGE_KEY] = {
"version": 1,
"minor_version": 1,
"data": {
"entities": [
{
"device_class": "best_class",
"entity_id": "test.entity",
"platform": "super_platform",
"unique_id": "very_unique",
},
]
},
}
await er.async_load(hass)
registry = er.async_get(hass)
entry = registry.async_get_or_create("test", "super_platform", "very_unique")
assert entry.device_class is None
assert entry.original_device_class == "best_class"
@pytest.mark.parametrize("load_registries", [False])
async def test_loading_invalid_entity_id(hass, hass_storage):
"""Test we skip entities with invalid entity IDs."""
@ -599,7 +632,7 @@ async def test_restore_states(hass):
suggested_object_id="all_info_set",
capabilities={"max": 100},
supported_features=5,
device_class="mock-device-class",
original_device_class="mock-device-class",
original_name="Mock Original Name",
original_icon="hass:original-icon",
)
@ -649,14 +682,14 @@ async def test_async_get_device_class_lookup(hass):
"light",
"battery_charging",
device_id="light_device_entry_id",
device_class="battery_charging",
original_device_class="battery_charging",
)
ent_reg.async_get_or_create(
"sensor",
"light",
"battery",
device_id="light_device_entry_id",
device_class="battery",
original_device_class="battery",
)
ent_reg.async_get_or_create(
"light", "light", "demo", device_id="light_device_entry_id"
@ -666,14 +699,14 @@ async def test_async_get_device_class_lookup(hass):
"vacuum",
"battery_charging",
device_id="vacuum_device_entry_id",
device_class="battery_charging",
original_device_class="battery_charging",
)
ent_reg.async_get_or_create(
"sensor",
"vacuum",
"battery",
device_id="vacuum_device_entry_id",
device_class="battery",
original_device_class="battery",
)
ent_reg.async_get_or_create(
"vacuum", "vacuum", "demo", device_id="vacuum_device_entry_id"
@ -683,7 +716,7 @@ async def test_async_get_device_class_lookup(hass):
"remote",
"battery_charging",
device_id="remote_device_entry_id",
device_class="battery_charging",
original_device_class="battery_charging",
)
ent_reg.async_get_or_create(
"remote", "remote", "demo", device_id="remote_device_entry_id"