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