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)
collection.attach_entity_component_collection(ent_comp, coll, MockEntity)
await coll.notify_change(
collection.CHANGE_ADDED,
"mock_id",
{"id": "mock_id", "state": "initial", "name": "Mock 1"},
)
assert hass.states.get("test.mock_1").name == "Mock 1"
assert hass.states.get("test.mock_1").state == "initial"
await coll.notify_change(
collection.CHANGE_UPDATED,
"mock_id",
{"id": "mock_id", "state": "second", "name": "Mock 1 updated"},
)
assert hass.states.get("test.mock_1").name == "Mock 1 updated"
assert hass.states.get("test.mock_1").state == "second"
await coll.notify_change(collection.CHANGE_REMOVED, "mock_id", None)
assert hass.states.get("test.mock_1") is None
async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
"""Set up configured zones as well as Home Assistant zone if necessary."""
component = entity_component.EntityComponent(_LOGGER, DOMAIN, hass)
id_manager = collection.IDManager()
yaml_collection = collection.IDLessCollection(
logging.getLogger(f"{__name__}.yaml_collection"), id_manager)
collection.sync_entity_lifecycle(hass, DOMAIN, DOMAIN, component,
yaml_collection, Zone.from_yaml)
storage_collection = ZoneStorageCollection(
storage.Store(hass, STORAGE_VERSION, STORAGE_KEY),
logging.getLogger(f"{__name__}.storage_collection"),
id_manager,
)
collection.sync_entity_lifecycle(hass, DOMAIN, DOMAIN, component,
storage_collection, Zone)
if config[DOMAIN]:
# AIS dom config can be empty
if config[DOMAIN] != [{}]:
await yaml_collection.async_load(config[DOMAIN])
await storage_collection.async_load()
collection.StorageCollectionWebsocket(storage_collection, DOMAIN, DOMAIN,
CREATE_FIELDS,
UPDATE_FIELDS).async_setup(hass)
async def reload_service_handler(service_call: ServiceCall) -> None:
"""Remove all zones and load new ones from config."""
conf = await component.async_prepare_reload(skip_reset=True)
if conf is None:
return
await yaml_collection.async_load(conf[DOMAIN])
service.async_register_admin_service(
hass,
DOMAIN,
SERVICE_RELOAD,
reload_service_handler,
schema=RELOAD_SERVICE_SCHEMA,
)
if component.get_entity("zone.home"):
return True
home_zone = Zone(_home_conf(hass))
home_zone.entity_id = ENTITY_ID_HOME
await component.async_add_entities([home_zone])
async def core_config_updated(_: Event) -> None:
"""Handle core config updated."""
await home_zone.async_update_config(_home_conf(hass))
hass.bus.async_listen(EVENT_CORE_CONFIG_UPDATE, core_config_updated)
hass.data[DOMAIN] = storage_collection
return True
async def async_setup(hass: HomeAssistant, config: Dict) -> bool:
"""Set up configured zones as well as Home Assistant zone if necessary."""
component = entity_component.EntityComponent(_LOGGER, DOMAIN, hass)
id_manager = collection.IDManager()
yaml_collection = collection.IDLessCollection(
logging.getLogger(f"{__name__}.yaml_collection"), id_manager
)
collection.attach_entity_component_collection(
component, yaml_collection, lambda conf: Zone(conf, False)
)
storage_collection = ZoneStorageCollection(
storage.Store(hass, STORAGE_VERSION, STORAGE_KEY),
logging.getLogger(f"{__name__}.storage_collection"),
id_manager,
)
collection.attach_entity_component_collection(
component, storage_collection, lambda conf: Zone(conf, True)
)
if config[DOMAIN]:
# AIS dom config can be empty
if config[DOMAIN] != [{}]:
await yaml_collection.async_load(config[DOMAIN])
await storage_collection.async_load()
collection.StorageCollectionWebsocket(
storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS
).async_setup(hass)
async def _collection_changed(change_type: str, item_id: str, config: Dict) -> None:
"""Handle a collection change: clean up entity registry on removals."""
if change_type != collection.CHANGE_REMOVED:
return
ent_reg = await entity_registry.async_get_registry(hass)
ent_reg.async_remove(
cast(str, ent_reg.async_get_entity_id(DOMAIN, DOMAIN, item_id))
)
storage_collection.async_add_listener(_collection_changed)
async def reload_service_handler(service_call: ServiceCall) -> None:
"""Remove all zones and load new ones from config."""
conf = await component.async_prepare_reload(skip_reset=True)
if conf is None:
return
await yaml_collection.async_load(conf[DOMAIN])
service.async_register_admin_service(
hass,
DOMAIN,
SERVICE_RELOAD,
reload_service_handler,
schema=RELOAD_SERVICE_SCHEMA,
)
if component.get_entity("zone.home"):
return True
home_zone = Zone(
_home_conf(hass),
True,
)
home_zone.entity_id = ENTITY_ID_HOME
await component.async_add_entities([home_zone])
async def core_config_updated(_: Event) -> None:
"""Handle core config updated."""
await home_zone.async_update_config(_home_conf(hass))
hass.bus.async_listen(EVENT_CORE_CONFIG_UPDATE, core_config_updated)
hass.data[DOMAIN] = storage_collection
return True
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)
async_update_config_calls = []
async_remove_calls = []
class MockMockEntity(MockEntity):
"""Track calls to async_update_config and async_remove."""
async def async_update_config(self, config):
nonlocal async_update_config_calls
async_update_config_calls.append(None)
await super().async_update_config()
async def async_remove(self, *, force_remove: bool = False):
nonlocal async_remove_calls
async_remove_calls.append(None)
await super().async_remove()
collection.sync_entity_lifecycle(
hass, "test", "test", ent_comp, coll, MockMockEntity
)
entity_registry = er.async_get(hass)
entity_registry.async_get_or_create(
"test", "test", "mock_id", suggested_object_id="mock_1"
)
await coll.notify_changes(
[
collection.CollectionChangeSet(
collection.CHANGE_ADDED,
"mock_id",
{"id": "mock_id", "state": "initial", "name": "Mock 1"},
)
],
)
assert hass.states.get("test.mock_1").name == "Mock 1"
assert hass.states.get("test.mock_1").state == "initial"
entity_registry.async_remove("test.mock_1")
await hass.async_block_till_done()
assert hass.states.get("test.mock_1") is None
assert len(async_remove_calls) == 1
await coll.notify_changes(
[
collection.CollectionChangeSet(
collection.CHANGE_UPDATED,
"mock_id",
{"id": "mock_id", "state": "second", "name": "Mock 1 updated"},
)
],
)
assert hass.states.get("test.mock_1") is None
assert len(async_update_config_calls) == 0
await coll.notify_changes(
[collection.CollectionChangeSet(collection.CHANGE_REMOVED, "mock_id", None)],
)
assert hass.states.get("test.mock_1") is None
assert len(async_remove_calls) == 1
