Home-Assistant
/
core
mirror of https://github.com/home-assistant/core.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add mixin class CollectionEntity for the collection helper (#77703) 

* Add mixin class CollectionEntity for the collection helper

* Improve typing

* Address review comments

* Fix tests
pull/77751/head
Erik Montnemery 2022-09-03 12:56:49 +02:00 committed by GitHub
parent 56278a4421
commit b0d033ef29
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
14 changed files with 221 additions and 86 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
homeassistant/components/counter/__init__.py
View File

@ -110,7 +110,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
logging.getLogger(f"{__name__}.yaml_collection"), id_manager
)
collection.sync_entity_lifecycle(
hass, DOMAIN, DOMAIN, component, yaml_collection, Counter.from_yaml
hass, DOMAIN, DOMAIN, component, yaml_collection, Counter
)
storage_collection = CounterStorageCollection(
@ -170,19 +170,26 @@ class CounterStorageCollection(collection.StorageCollection):
return {**data, **update_data}
class Counter(RestoreEntity):
class Counter(collection.CollectionEntity, RestoreEntity):
"""Representation of a counter."""
_attr_should_poll: bool = False
editable: bool
def __init__(self, config: dict) -> None:
def __init__(self, config: ConfigType) -> None:
"""Initialize a counter."""
self._config: dict = config
self._config: ConfigType = config
self._state: int | None = config[CONF_INITIAL]
self.editable: bool = True
@classmethod
def from_yaml(cls, config: dict) -> Counter:
def from_storage(cls, config: ConfigType) -> Counter:
"""Create counter instance from storage."""
counter = cls(config)
counter.editable = True
return counter
@classmethod
def from_yaml(cls, config: ConfigType) -> Counter:
"""Create counter instance from yaml config."""
counter = cls(config)
counter.editable = False
@ -273,7 +280,7 @@ class Counter(RestoreEntity):
self._state = self.compute_next_state(new_state)
self.async_write_ha_state()
async def async_update_config(self, config: dict) -> None:
async def async_update_config(self, config: ConfigType) -> None:
"""Change the counter's settings WS CRUD."""
self._config = config
self._state = self.compute_next_state(self._state)

15
homeassistant/components/input_boolean/__init__.py
View File

@ -100,7 +100,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
logging.getLogger(f"{__name__}.yaml_collection"), id_manager
)
collection.sync_entity_lifecycle(
hass, DOMAIN, DOMAIN, component, yaml_collection, InputBoolean.from_yaml
hass, DOMAIN, DOMAIN, component, yaml_collection, InputBoolean
)
storage_collection = InputBooleanStorageCollection(
@ -150,21 +150,28 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
return True
class InputBoolean(ToggleEntity, RestoreEntity):
class InputBoolean(collection.CollectionEntity, ToggleEntity, RestoreEntity):
"""Representation of a boolean input."""
_attr_should_poll = False
editable: bool
def __init__(self, config: ConfigType) -> None:
"""Initialize a boolean input."""
self._config = config
self.editable = True
self._attr_is_on = config.get(CONF_INITIAL, False)
self._attr_unique_id = config[CONF_ID]
@classmethod
def from_storage(cls, config: ConfigType) -> InputBoolean:
"""Return entity instance initialized from storage."""
input_bool = cls(config)
input_bool.editable = True
return input_bool
@classmethod
def from_yaml(cls, config: ConfigType) -> InputBoolean:
"""Return entity instance initialized from yaml storage."""
"""Return entity instance initialized from yaml."""
input_bool = cls(config)
input_bool.entity_id = f"{DOMAIN}.{config[CONF_ID]}"
input_bool.editable = False

17
homeassistant/components/input_button/__init__.py
View File

@ -85,7 +85,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
logging.getLogger(f"{__name__}.yaml_collection"), id_manager
)
collection.sync_entity_lifecycle(
hass, DOMAIN, DOMAIN, component, yaml_collection, InputButton.from_yaml
hass, DOMAIN, DOMAIN, component, yaml_collection, InputButton
)
storage_collection = InputButtonStorageCollection(
@ -131,20 +131,27 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
return True
class InputButton(ButtonEntity, RestoreEntity):
class InputButton(collection.CollectionEntity, ButtonEntity, RestoreEntity):
"""Representation of a button."""
_attr_should_poll = False
editable: bool
def __init__(self, config: ConfigType) -> None:
"""Initialize a button."""
self._config = config
self.editable = True
self._attr_unique_id = config[CONF_ID]
@classmethod
def from_yaml(cls, config: ConfigType) -> ButtonEntity:
"""Return entity instance initialized from yaml storage."""
def from_storage(cls, config: ConfigType) -> InputButton:
"""Return entity instance initialized from storage."""
button = cls(config)
button.editable = True
return button
@classmethod
def from_yaml(cls, config: ConfigType) -> InputButton:
"""Return entity instance initialized from yaml."""
button = cls(config)
button.entity_id = f"{DOMAIN}.{config[CONF_ID]}"
button.editable = False

21
homeassistant/components/input_datetime/__init__.py
View File

@ -149,7 +149,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
logging.getLogger(f"{__name__}.yaml_collection"), id_manager
)
collection.sync_entity_lifecycle(
hass, DOMAIN, DOMAIN, component, yaml_collection, InputDatetime.from_yaml
hass, DOMAIN, DOMAIN, component, yaml_collection, InputDatetime
)
storage_collection = DateTimeStorageCollection(
@ -231,15 +231,15 @@ class DateTimeStorageCollection(collection.StorageCollection):
return has_date_or_time({**data, **update_data})
class InputDatetime(RestoreEntity):
class InputDatetime(collection.CollectionEntity, RestoreEntity):
"""Representation of a datetime input."""
_attr_should_poll = False
editable: bool
def __init__(self, config: dict) -> None:
def __init__(self, config: ConfigType) -> None:
"""Initialize a select input."""
self._config = config
self.editable = True
self._current_datetime = None
if not config.get(CONF_INITIAL):
@ -258,8 +258,15 @@ class InputDatetime(RestoreEntity):
)
@classmethod
def from_yaml(cls, config: dict) -> InputDatetime:
"""Return entity instance initialized from yaml storage."""
def from_storage(cls, config: ConfigType) -> InputDatetime:
"""Return entity instance initialized from storage."""
input_dt = cls(config)
input_dt.editable = True
return input_dt
@classmethod
def from_yaml(cls, config: ConfigType) -> InputDatetime:
"""Return entity instance initialized from yaml."""
input_dt = cls(config)
input_dt.entity_id = f"{DOMAIN}.{config[CONF_ID]}"
input_dt.editable = False
@ -420,7 +427,7 @@ class InputDatetime(RestoreEntity):
)
self.async_write_ha_state()
async def async_update_config(self, config: dict) -> None:
async def async_update_config(self, config: ConfigType) -> None:
"""Handle when the config is updated."""
self._config = config
self.async_write_ha_state()

21
homeassistant/components/input_number/__init__.py
View File

@ -130,7 +130,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
logging.getLogger(f"{__name__}.yaml_collection"), id_manager
)
collection.sync_entity_lifecycle(
hass, DOMAIN, DOMAIN, component, yaml_collection, InputNumber.from_yaml
hass, DOMAIN, DOMAIN, component, yaml_collection, InputNumber
)
storage_collection = NumberStorageCollection(
@ -202,20 +202,27 @@ class NumberStorageCollection(collection.StorageCollection):
return _cv_input_number({**data, **update_data})
class InputNumber(RestoreEntity):
class InputNumber(collection.CollectionEntity, RestoreEntity):
"""Representation of a slider."""
_attr_should_poll = False
editable: bool
def __init__(self, config: dict) -> None:
def __init__(self, config: ConfigType) -> None:
"""Initialize an input number."""
self._config = config
self.editable = True
self._current_value: float | None = config.get(CONF_INITIAL)
@classmethod
def from_yaml(cls, config: dict) -> InputNumber:
"""Return entity instance initialized from yaml storage."""
def from_storage(cls, config: ConfigType) -> InputNumber:
"""Return entity instance initialized from storage."""
input_num = cls(config)
input_num.editable = True
return input_num
@classmethod
def from_yaml(cls, config: ConfigType) -> InputNumber:
"""Return entity instance initialized from yaml."""
input_num = cls(config)
input_num.entity_id = f"{DOMAIN}.{config[CONF_ID]}"
input_num.editable = False
@ -310,7 +317,7 @@ class InputNumber(RestoreEntity):
"""Decrement value."""
await self.async_set_value(max(self._current_value - self._step, self._minimum))
async def async_update_config(self, config: dict) -> None:
async def async_update_config(self, config: ConfigType) -> None:
"""Handle when the config is updated."""
self._config = config
# just in case min/max values changed

15
homeassistant/components/input_select/__init__.py
View File

@ -152,7 +152,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
logging.getLogger(f"{__name__}.yaml_collection"), id_manager
)
collection.sync_entity_lifecycle(
hass, DOMAIN, DOMAIN, component, yaml_collection, InputSelect.from_yaml
hass, DOMAIN, DOMAIN, component, yaml_collection, InputSelect
)
storage_collection = InputSelectStorageCollection(
@ -258,11 +258,11 @@ class InputSelectStorageCollection(collection.StorageCollection):
return _cv_input_select({**data, **update_data})
class InputSelect(SelectEntity, RestoreEntity):
class InputSelect(collection.CollectionEntity, SelectEntity, RestoreEntity):
"""Representation of a select input."""
_attr_should_poll = False
editable = True
editable: bool
def __init__(self, config: ConfigType) -> None:
"""Initialize a select input."""
@ -272,9 +272,16 @@ class InputSelect(SelectEntity, RestoreEntity):
self._attr_options = config[CONF_OPTIONS]
self._attr_unique_id = config[CONF_ID]
@classmethod
def from_storage(cls, config: ConfigType) -> InputSelect:
"""Return entity instance initialized from storage."""
input_select = cls(config)
input_select.editable = True
return input_select
@classmethod
def from_yaml(cls, config: ConfigType) -> InputSelect:
"""Return entity instance initialized from yaml storage."""
"""Return entity instance initialized from yaml."""
input_select = cls(config)
input_select.entity_id = f"{DOMAIN}.{config[CONF_ID]}"
input_select.editable = False

21
homeassistant/components/input_text/__init__.py
View File

@ -129,7 +129,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
logging.getLogger(f"{__name__}.yaml_collection"), id_manager
)
collection.sync_entity_lifecycle(
hass, DOMAIN, DOMAIN, component, yaml_collection, InputText.from_yaml
hass, DOMAIN, DOMAIN, component, yaml_collection, InputText
)
storage_collection = InputTextStorageCollection(
@ -195,20 +195,27 @@ class InputTextStorageCollection(collection.StorageCollection):
return _cv_input_text({**data, **update_data})
class InputText(RestoreEntity):
class InputText(collection.CollectionEntity, RestoreEntity):
"""Represent a text box."""
_attr_should_poll = False
editable: bool
def __init__(self, config: dict) -> None:
def __init__(self, config: ConfigType) -> None:
"""Initialize a text input."""
self._config = config
self.editable = True
self._current_value = config.get(CONF_INITIAL)
@classmethod
def from_yaml(cls, config: dict) -> InputText:
"""Return entity instance initialized from yaml storage."""
def from_storage(cls, config: ConfigType) -> InputText:
"""Return entity instance initialized from storage."""
input_text = cls(config)
input_text.editable = True
return input_text
@classmethod
def from_yaml(cls, config: ConfigType) -> InputText:
"""Return entity instance initialized from yaml."""
input_text = cls(config)
input_text.entity_id = f"{DOMAIN}.{config[CONF_ID]}"
input_text.editable = False
@ -286,7 +293,7 @@ class InputText(RestoreEntity):
self._current_value = value
self.async_write_ha_state()
async def async_update_config(self, config: dict) -> None:
async def async_update_config(self, config: ConfigType) -> None:
"""Handle when the config is updated."""
self._config = config
self.async_write_ha_state()

19
homeassistant/components/person/__init__.py
View File

@ -342,7 +342,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
hass, DOMAIN, DOMAIN, entity_component, yaml_collection, Person
)
collection.sync_entity_lifecycle(
hass, DOMAIN, DOMAIN, entity_component, storage_collection, Person.from_yaml
hass, DOMAIN, DOMAIN, entity_component, storage_collection, Person
)
await yaml_collection.async_load(
@ -385,15 +385,15 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
return True
class Person(RestoreEntity):
class Person(collection.CollectionEntity, RestoreEntity):
"""Represent a tracked person."""
_attr_should_poll = False
editable: bool
def __init__(self, config):
"""Set up person."""
self._config = config
self.editable = True
self._latitude = None
self._longitude = None
self._gps_accuracy = None
@ -402,8 +402,15 @@ class Person(RestoreEntity):
self._unsub_track_device = None
@classmethod
def from_yaml(cls, config):
"""Return entity instance initialized from yaml storage."""
def from_storage(cls, config: ConfigType):
"""Return entity instance initialized from storage."""
person = cls(config)
person.editable = True
return person
@classmethod
def from_yaml(cls, config: ConfigType):
"""Return entity instance initialized from yaml."""
person = cls(config)
person.editable = False
return person
@ -468,7 +475,7 @@ class Person(RestoreEntity):
EVENT_HOMEASSISTANT_START, person_start_hass
)
async def async_update_config(self, config):
async def async_update_config(self, config: ConfigType):
"""Handle when the config is updated."""
self._config = config

18
homeassistant/components/schedule/__init__.py
View File

@ -20,6 +20,7 @@ from homeassistant.const import (
)
from homeassistant.core import HomeAssistant, ServiceCall, callback
from homeassistant.helpers.collection import (
CollectionEntity,
IDManager,
StorageCollection,
StorageCollectionWebsocket,
@ -27,7 +28,6 @@ from homeassistant.helpers.collection import (
sync_entity_lifecycle,
)
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.entity import Entity
from homeassistant.helpers.entity_component import EntityComponent
from homeassistant.helpers.event import async_track_point_in_utc_time
from homeassistant.helpers.integration_platform import (
@ -163,9 +163,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
id_manager = IDManager()
yaml_collection = YamlCollection(LOGGER, id_manager)
sync_entity_lifecycle(
hass, DOMAIN, DOMAIN, component, yaml_collection, Schedule.from_yaml
)
sync_entity_lifecycle(hass, DOMAIN, DOMAIN, component, yaml_collection, Schedule)
storage_collection = ScheduleStorageCollection(
Store(
@ -239,7 +237,7 @@ class ScheduleStorageCollection(StorageCollection):
return data
class Schedule(Entity):
class Schedule(CollectionEntity):
"""Schedule entity."""
_attr_has_entity_name = True
@ -249,7 +247,7 @@ class Schedule(Entity):
_next: datetime
_unsub_update: Callable[[], None] | None = None
def __init__(self, config: ConfigType, editable: bool = True) -> None:
def __init__(self, config: ConfigType, editable: bool) -> None:
"""Initialize a schedule."""
self._config = ENTITY_SCHEMA(config)
self._attr_capability_attributes = {ATTR_EDITABLE: editable}
@ -257,9 +255,15 @@ class Schedule(Entity):
self._attr_name = self._config[CONF_NAME]
self._attr_unique_id = self._config[CONF_ID]
@classmethod
def from_storage(cls, config: ConfigType) -> Schedule:
"""Return entity instance initialized from storage."""
schedule = cls(config, editable=True)
return schedule
@classmethod
def from_yaml(cls, config: ConfigType) -> Schedule:
"""Return entity instance initialized from yaml storage."""
"""Return entity instance initialized from yaml."""
schedule = cls(config, editable=False)
schedule.entity_id = f"{DOMAIN}.{config[CONF_ID]}"
return schedule

22
homeassistant/components/timer/__init__.py
View File

@ -119,7 +119,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
logging.getLogger(f"{__name__}.yaml_collection"), id_manager
)
collection.sync_entity_lifecycle(
hass, DOMAIN, DOMAIN, component, yaml_collection, Timer.from_yaml
hass, DOMAIN, DOMAIN, component, yaml_collection, Timer
)
storage_collection = TimerStorageCollection(
@ -195,13 +195,14 @@ class TimerStorageCollection(collection.StorageCollection):
return data
class Timer(RestoreEntity):
class Timer(collection.CollectionEntity, RestoreEntity):
"""Representation of a timer."""
def __init__(self, config: dict) -> None:
editable: bool
def __init__(self, config: ConfigType) -> None:
"""Initialize a timer."""
self._config: dict = config
self.editable: bool = True
self._state: str = STATUS_IDLE
self._duration = cv.time_period_str(config[CONF_DURATION])
self._remaining: timedelta | None = None
@ -213,8 +214,15 @@ class Timer(RestoreEntity):
self._attr_force_update = True
@classmethod
def from_yaml(cls, config: dict) -> Timer:
"""Return entity instance initialized from yaml storage."""
def from_storage(cls, config: ConfigType) -> Timer:
"""Return entity instance initialized from storage."""
timer = cls(config)
timer.editable = True
return timer
@classmethod
def from_yaml(cls, config: ConfigType) -> Timer:
"""Return entity instance initialized from yaml."""
timer = cls(config)
timer.entity_id = ENTITY_ID_FORMAT.format(config[CONF_ID])
timer.editable = False
@ -384,7 +392,7 @@ class Timer(RestoreEntity):
)
self.async_write_ha_state()
async def async_update_config(self, config: dict) -> None:
async def async_update_config(self, config: ConfigType) -> None:
"""Handle when the config is updated."""
self._config = config
self._duration = cv.time_period_str(config[CONF_DURATION])

24
homeassistant/components/zone/__init__.py
View File

@ -31,7 +31,6 @@ from homeassistant.core import Event, HomeAssistant, ServiceCall, State, callbac
from homeassistant.helpers import (
collection,
config_validation as cv,
entity,
entity_component,
event,
service,
@ -193,7 +192,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
logging.getLogger(f"{__name__}.yaml_collection"), id_manager
)
collection.sync_entity_lifecycle(
hass, DOMAIN, DOMAIN, component, yaml_collection, Zone.from_yaml
hass, DOMAIN, DOMAIN, component, yaml_collection, Zone
)
storage_collection = ZoneStorageCollection(
@ -284,21 +283,30 @@ async def async_unload_entry(
return True
class Zone(entity.Entity):
class Zone(collection.CollectionEntity):
"""Representation of a Zone."""
def __init__(self, config: dict) -> None:
editable: bool
def __init__(self, config: ConfigType) -> None:
"""Initialize the zone."""
self._config = config
self.editable = True
self._attrs: dict | None = None
self._remove_listener: Callable[[], None] | None = None
self._persons_in_zone: set[str] = set()
self._generate_attrs()
@classmethod
def from_yaml(cls, config: dict) -> Zone:
"""Return entity instance initialized from yaml storage."""
def from_storage(cls, config: ConfigType) -> Zone:
"""Return entity instance initialized from storage."""
zone = cls(config)
zone.editable = True
zone._generate_attrs()
return zone
@classmethod
def from_yaml(cls, config: ConfigType) -> Zone:
"""Return entity instance initialized from yaml."""
zone = cls(config)
zone.editable = False
zone._generate_attrs()
@ -329,7 +337,7 @@ class Zone(entity.Entity):
"""Zone does not poll."""
return False
async def async_update_config(self, config: dict) -> None:
async def async_update_config(self, config: ConfigType) -> None:
"""Handle when the config is updated."""
if self._config == config:
return

52
homeassistant/helpers/collection.py
View File

@ -22,6 +22,7 @@ from . import entity_registry
from .entity import Entity
from .entity_component import EntityComponent
from .storage import Store
from .typing import ConfigType
STORAGE_VERSION = 1
SAVE_DELAY = 10
@ -101,6 +102,24 @@ class IDManager:
return proposal
class CollectionEntity(Entity):
"""Mixin class for entities managed by an ObservableCollection."""
@classmethod
@abstractmethod
def from_storage(cls, config: ConfigType) -> CollectionEntity:
"""Create instance from storage."""
@classmethod
@abstractmethod
def from_yaml(cls, config: ConfigType) -> CollectionEntity:
"""Create instance from yaml config."""
@abstractmethod
async def async_update_config(self, config: ConfigType) -> None:
"""Handle updated configuration."""
class ObservableCollection(ABC):
"""Base collection type that can be observed."""
@ -155,6 +174,13 @@ class ObservableCollection(ABC):
class YamlCollection(ObservableCollection):
"""Offer a collection based on static data."""
@staticmethod
def create_entity(
entity_class: type[CollectionEntity], config: ConfigType
) -> CollectionEntity:
"""Create a CollectionEntity instance."""
return entity_class.from_yaml(config)
async def async_load(self, data: list[dict]) -> None:
"""Load the YAML collection. Overrides existing data."""
old_ids = set(self.data)
@ -198,6 +224,13 @@ class StorageCollection(ObservableCollection):
super().__init__(logger, id_manager)
self.store = store
@staticmethod
def create_entity(
entity_class: type[CollectionEntity], config: ConfigType
) -> CollectionEntity:
"""Create a CollectionEntity instance."""
return entity_class.from_storage(config)
@property
def hass(self) -> HomeAssistant:
"""Home Assistant object."""
@ -290,7 +323,7 @@ class StorageCollection(ObservableCollection):
return {"items": list(self.data.values())}
class IDLessCollection(ObservableCollection):
class IDLessCollection(YamlCollection):
"""A collection without IDs."""
counter = 0
@ -326,20 +359,22 @@ def sync_entity_lifecycle(
domain: str,
platform: str,
entity_component: EntityComponent,
collection: ObservableCollection,
create_entity: Callable[[dict], Entity],
collection: StorageCollection | YamlCollection,
entity_class: type[CollectionEntity],
) -> None:
"""Map a collection to an entity component."""
entities: dict[str, Entity] = {}
entities: dict[str, CollectionEntity] = {}
ent_reg = entity_registry.async_get(hass)
async def _add_entity(change_set: CollectionChangeSet) -> Entity:
async def _add_entity(change_set: CollectionChangeSet) -> CollectionEntity:
def entity_removed() -> None:
"""Remove entity from entities if it's removed or not added."""
if change_set.item_id in entities:
entities.pop(change_set.item_id)
entities[change_set.item_id] = create_entity(change_set.item)
entities[change_set.item_id] = collection.create_entity(
entity_class, change_set.item
)
entities[change_set.item_id].async_on_remove(entity_removed)
return entities[change_set.item_id]
@ -359,10 +394,11 @@ def sync_entity_lifecycle(
async def _update_entity(change_set: CollectionChangeSet) -> None:
if change_set.item_id not in entities:
return
await entities[change_set.item_id].async_update_config(change_set.item) # type: ignore[attr-defined]
await entities[change_set.item_id].async_update_config(change_set.item)
_func_map: dict[
str, Callable[[CollectionChangeSet], Coroutine[Any, Any, Entity | None]]
str,
Callable[[CollectionChangeSet], Coroutine[Any, Any, CollectionEntity | None]],
] = {
CHANGE_ADDED: _add_entity,
CHANGE_REMOVED: _remove_entity,

8
tests/components/timer/test_init.py
View File

@ -672,7 +672,7 @@ async def test_restore_idle(hass):
# Emulate a fresh load
hass.data.pop(DATA_RESTORE_STATE_TASK)
entity = Timer(
entity = Timer.from_storage(
{
CONF_ID: "test",
CONF_NAME: "test",
@ -712,7 +712,7 @@ async def test_restore_paused(hass):
# Emulate a fresh load
hass.data.pop(DATA_RESTORE_STATE_TASK)
entity = Timer(
entity = Timer.from_storage(
{
CONF_ID: "test",
CONF_NAME: "test",
@ -756,7 +756,7 @@ async def test_restore_active_resume(hass):
# Emulate a fresh load
hass.data.pop(DATA_RESTORE_STATE_TASK)
entity = Timer(
entity = Timer.from_storage(
{
CONF_ID: "test",
CONF_NAME: "test",
@ -807,7 +807,7 @@ async def test_restore_active_finished_outside_grace(hass):
# Emulate a fresh load
hass.data.pop(DATA_RESTORE_STATE_TASK)
entity = Timer(
entity = Timer.from_storage(
{
CONF_ID: "test",
CONF_NAME: "test",

33
tests/helpers/test_collection.py
View File

@ -1,4 +1,6 @@
"""Tests for the collection helper."""
from __future__ import annotations
import logging
import pytest
@ -6,11 +8,11 @@ import voluptuous as vol
from homeassistant.helpers import (
collection,
entity,
entity_component,
entity_registry as er,
storage,
)
from homeassistant.helpers.typing import ConfigType
from tests.common import flush_store
@ -29,13 +31,23 @@ def track_changes(coll: collection.ObservableCollection):
return changes
class MockEntity(entity.Entity):
class MockEntity(collection.CollectionEntity):
"""Entity that is config based."""
def __init__(self, config):
"""Initialize entity."""
self._config = config
@classmethod
def from_storage(cls, config: ConfigType) -> MockEntity:
"""Create instance from storage."""
return cls(config)
@classmethod
def from_yaml(cls, config: ConfigType) -> MockEntity:
"""Create instance from storage."""
raise NotImplementedError
@property
def unique_id(self):
"""Return unique ID of entity."""
@ -57,6 +69,17 @@ class MockEntity(entity.Entity):
self.async_write_ha_state()
class MockObservableCollection(collection.ObservableCollection):
"""Mock observable collection which can create entities."""
@staticmethod
def create_entity(
entity_class: type[collection.CollectionEntity], config: ConfigType
) -> collection.CollectionEntity:
"""Create a CollectionEntity instance."""
return entity_class.from_storage(config)
class MockStorageCollection(collection.StorageCollection):
"""Mock storage collection."""
@ -231,7 +254,7 @@ async def test_storage_collection(hass):
async def test_attach_entity_component_collection(hass):
"""Test attaching collection to entity component."""
ent_comp = entity_component.EntityComponent(_LOGGER, "test", hass)
coll = collection.ObservableCollection(_LOGGER)
coll = MockObservableCollection(_LOGGER)
collection.sync_entity_lifecycle(hass, "test", "test", ent_comp, coll, MockEntity)
await coll.notify_changes(
@ -270,7 +293,7 @@ async def test_attach_entity_component_collection(hass):
async def test_entity_component_collection_abort(hass):
"""Test aborted entity adding is handled."""
ent_comp = entity_component.EntityComponent(_LOGGER, "test", hass)
coll = collection.ObservableCollection(_LOGGER)
coll = MockObservableCollection(_LOGGER)
async_update_config_calls = []
async_remove_calls = []
@ -336,7 +359,7 @@ async def test_entity_component_collection_abort(hass):
async def test_entity_component_collection_entity_removed(hass):
"""Test entity removal is handled."""
ent_comp = entity_component.EntityComponent(_LOGGER, "test", hass)
coll = collection.ObservableCollection(_LOGGER)
coll = MockObservableCollection(_LOGGER)
async_update_config_calls = []
async_remove_calls = []