Improve test

This commit is contained in:
Erik 2025-06-14 09:08:23 +02:00
parent b697bf13d2
commit 3cf2ed5898

View File

@ -814,29 +814,68 @@ async def test_warn_slow_write_state_custom_component(
) in caplog.text ) in caplog.text
async def test_setup_source(hass: HomeAssistant) -> None: async def test_setup_source(
hass: HomeAssistant, entity_registry: er.EntityRegistry
) -> None:
"""Check that we register sources correctly.""" """Check that we register sources correctly."""
platform = MockEntityPlatform(hass) platform = MockEntityPlatform(hass)
entity_platform = MockEntity(name="Platform Config Source") entity_platform = MockEntity(name="Platform Config Source 1")
await platform.async_add_entities([entity_platform]) entity_platform_uniq = MockEntity(name="Platform Config Source 2", unique_id="ABC")
await platform.async_add_entities([entity_platform, entity_platform_uniq])
platform.config_entry = MockConfigEntry() platform.config_entry = MockConfigEntry()
entity_entry = MockEntity(name="Config Entry Source") platform.config_entry.add_to_hass(hass)
await platform.async_add_entities([entity_entry]) entity_entry = MockEntity(name="Config Entry Source 1")
entity_entry_uniq = MockEntity(name="Config Entry Source 2", unique_id="DEF")
await platform.async_add_entities([entity_entry, entity_entry_uniq])
assert {state.entity_id: state.state for state in hass.states.async_all()} == {
"test_domain.config_entry_source_1": "unknown",
"test_domain.config_entry_source_2": "unknown",
"test_domain.platform_config_source_1": "unknown",
"test_domain.platform_config_source_2": "unknown",
}
assert entity.entity_sources(hass) == { assert entity.entity_sources(hass) == {
"test_domain.platform_config_source": { "test_domain.platform_config_source_1": {
"domain": "test_platform", "domain": "test_platform",
}, },
"test_domain.config_entry_source": { "test_domain.platform_config_source_2": {
"domain": "test_platform",
},
"test_domain.config_entry_source_1": {
"config_entry": platform.config_entry.entry_id,
"domain": "test_platform",
},
"test_domain.config_entry_source_2": {
"config_entry": platform.config_entry.entry_id, "config_entry": platform.config_entry.entry_id,
"domain": "test_platform", "domain": "test_platform",
}, },
} }
# Reset the platform to remove all the entities
await platform.async_reset() await platform.async_reset()
# Check that entities with unique IDs are still present
assert {state.entity_id: state.state for state in hass.states.async_all()} == {
"test_domain.config_entry_source_2": "unavailable",
"test_domain.platform_config_source_2": "unavailable",
}
assert entity.entity_sources(hass) == {
"test_domain.platform_config_source_2": {
"domain": "test_platform",
},
"test_domain.config_entry_source_2": {
"config_entry": platform.config_entry.entry_id,
"domain": "test_platform",
},
}
# Remove the entity registry entries for the entities with unique IDs
entity_registry.async_remove("test_domain.platform_config_source_2")
entity_registry.async_remove("test_domain.config_entry_source_2")
assert hass.states.async_all() == []
assert entity.entity_sources(hass) == {} assert entity.entity_sources(hass) == {}