From 5bb96f7f061757352c0a533d2f1d3fdaeec23eed Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Tue, 26 Aug 2025 12:37:56 +0200 Subject: [PATCH] Adjust device disabled_by flag when changing config entry (#151155) --- .../components/anthropic/__init__.py | 12 +- .../__init__.py | 12 +- homeassistant/components/ollama/__init__.py | 12 +- .../openai_conversation/__init__.py | 12 +- homeassistant/helpers/device_registry.py | 26 ++ tests/components/anthropic/test_init.py | 28 +- .../test_init.py | 28 +- tests/components/ollama/test_init.py | 28 +- .../openai_conversation/test_init.py | 28 +- tests/helpers/test_device_registry.py | 260 ++++++++++++++++++ 10 files changed, 390 insertions(+), 56 deletions(-) diff --git a/homeassistant/components/anthropic/__init__.py b/homeassistant/components/anthropic/__init__.py index b996b7d38c5..55178d101fb 100644 --- a/homeassistant/components/anthropic/__init__.py +++ b/homeassistant/components/anthropic/__init__.py @@ -129,9 +129,9 @@ async def async_migrate_integration(hass: HomeAssistant) -> None: entity_disabled_by is er.RegistryEntryDisabler.CONFIG_ENTRY and not all_disabled ): - # Device and entity registries don't update the disabled_by flag - # when moving a device or entity from one config entry to another, - # so we need to do it manually. + # Device and entity registries will set the disabled_by flag to None + # when moving a device or entity disabled by CONFIG_ENTRY to an enabled + # config entry, but we want to set it to DEVICE or USER instead, entity_disabled_by = ( er.RegistryEntryDisabler.DEVICE if device @@ -146,9 +146,9 @@ async def async_migrate_integration(hass: HomeAssistant) -> None: ) if device is not None: - # Device and entity registries don't update the disabled_by flag when - # moving a device or entity from one config entry to another, so we - # need to do it manually. + # Device and entity registries will set the disabled_by flag to None + # when moving a device or entity disabled by CONFIG_ENTRY to an enabled + # config entry, but we want to set it to USER instead, device_disabled_by = device.disabled_by if ( device.disabled_by is dr.DeviceEntryDisabler.CONFIG_ENTRY diff --git a/homeassistant/components/google_generative_ai_conversation/__init__.py b/homeassistant/components/google_generative_ai_conversation/__init__.py index a1fd5ea0f9b..8d7fb1b1cc4 100644 --- a/homeassistant/components/google_generative_ai_conversation/__init__.py +++ b/homeassistant/components/google_generative_ai_conversation/__init__.py @@ -260,9 +260,9 @@ async def async_migrate_integration(hass: HomeAssistant) -> None: entity_disabled_by is er.RegistryEntryDisabler.CONFIG_ENTRY and not all_disabled ): - # Device and entity registries don't update the disabled_by flag - # when moving a device or entity from one config entry to another, - # so we need to do it manually. + # Device and entity registries will set the disabled_by flag to None + # when moving a device or entity disabled by CONFIG_ENTRY to an enabled + # config entry, but we want to set it to DEVICE or USER instead, entity_disabled_by = ( er.RegistryEntryDisabler.DEVICE if device @@ -277,9 +277,9 @@ async def async_migrate_integration(hass: HomeAssistant) -> None: ) if device is not None: - # Device and entity registries don't update the disabled_by flag when - # moving a device or entity from one config entry to another, so we - # need to do it manually. + # Device and entity registries will set the disabled_by flag to None + # when moving a device or entity disabled by CONFIG_ENTRY to an enabled + # config entry, but we want to set it to USER instead, device_disabled_by = device.disabled_by if ( device.disabled_by is dr.DeviceEntryDisabler.CONFIG_ENTRY diff --git a/homeassistant/components/ollama/__init__.py b/homeassistant/components/ollama/__init__.py index 091e58dbe7f..805724b82e3 100644 --- a/homeassistant/components/ollama/__init__.py +++ b/homeassistant/components/ollama/__init__.py @@ -145,9 +145,9 @@ async def async_migrate_integration(hass: HomeAssistant) -> None: entity_disabled_by is er.RegistryEntryDisabler.CONFIG_ENTRY and not all_disabled ): - # Device and entity registries don't update the disabled_by flag - # when moving a device or entity from one config entry to another, - # so we need to do it manually. + # Device and entity registries will set the disabled_by flag to None + # when moving a device or entity disabled by CONFIG_ENTRY to an enabled + # config entry, but we want to set it to DEVICE or USER instead, entity_disabled_by = ( er.RegistryEntryDisabler.DEVICE if device @@ -162,9 +162,9 @@ async def async_migrate_integration(hass: HomeAssistant) -> None: ) if device is not None: - # Device and entity registries don't update the disabled_by flag when - # moving a device or entity from one config entry to another, so we - # need to do it manually. + # Device and entity registries will set the disabled_by flag to None + # when moving a device or entity disabled by CONFIG_ENTRY to an enabled + # config entry, but we want to set it to USER instead, device_disabled_by = device.disabled_by if ( device.disabled_by is dr.DeviceEntryDisabler.CONFIG_ENTRY diff --git a/homeassistant/components/openai_conversation/__init__.py b/homeassistant/components/openai_conversation/__init__.py index f50563b59ea..06a61d70b01 100644 --- a/homeassistant/components/openai_conversation/__init__.py +++ b/homeassistant/components/openai_conversation/__init__.py @@ -320,9 +320,9 @@ async def async_migrate_integration(hass: HomeAssistant) -> None: entity_disabled_by is er.RegistryEntryDisabler.CONFIG_ENTRY and not all_disabled ): - # Device and entity registries don't update the disabled_by flag - # when moving a device or entity from one config entry to another, - # so we need to do it manually. + # Device and entity registries will set the disabled_by flag to None + # when moving a device or entity disabled by CONFIG_ENTRY to an enabled + # config entry, but we want to set it to DEVICE or USER instead, entity_disabled_by = ( er.RegistryEntryDisabler.DEVICE if device @@ -337,9 +337,9 @@ async def async_migrate_integration(hass: HomeAssistant) -> None: ) if device is not None: - # Device and entity registries don't update the disabled_by flag when - # moving a device or entity from one config entry to another, so we - # need to do it manually. + # Device and entity registries will set the disabled_by flag to None + # when moving a device or entity disabled by CONFIG_ENTRY to an enabled + # config entry, but we want to set it to USER instead, device_disabled_by = device.disabled_by if ( device.disabled_by is dr.DeviceEntryDisabler.CONFIG_ENTRY diff --git a/homeassistant/helpers/device_registry.py b/homeassistant/helpers/device_registry.py index a78fa935606..e25ca11e083 100644 --- a/homeassistant/helpers/device_registry.py +++ b/homeassistant/helpers/device_registry.py @@ -1115,6 +1115,16 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]): config_entries_subentries = old.config_entries_subentries | { add_config_entry_id: {add_config_subentry_id} } + # Enable the device if it was disabled by config entry and we're adding + # a non disabled config entry + if ( + # mypy says add_config_entry can be None. That's impossible, because we + # raise above if that happens + not add_config_entry.disabled_by # type: ignore[union-attr] + and old.disabled_by is DeviceEntryDisabler.CONFIG_ENTRY + ): + new_values["disabled_by"] = None + old_values["disabled_by"] = old.disabled_by elif ( add_config_subentry_id not in old.config_entries_subentries[add_config_entry_id] @@ -1157,6 +1167,22 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]): config_entries = config_entries - {remove_config_entry_id} + # Disable the device if it is enabled and all remaining config entries + # are disabled + has_enabled_config_entries = any( + config_entry.disabled_by is None + for config_entry_id in config_entries + if ( + config_entry := self.hass.config_entries.async_get_entry( + config_entry_id + ) + ) + is not None + ) + if not has_enabled_config_entries and old.disabled_by is None: + new_values["disabled_by"] = DeviceEntryDisabler.CONFIG_ENTRY + old_values["disabled_by"] = old.disabled_by + if config_entries != old.config_entries: new_values["config_entries"] = config_entries old_values["config_entries"] = old.config_entries diff --git a/tests/components/anthropic/test_init.py b/tests/components/anthropic/test_init.py index ff54539bb39..a97a3b7a378 100644 --- a/tests/components/anthropic/test_init.py +++ b/tests/components/anthropic/test_init.py @@ -156,6 +156,8 @@ async def test_migration_from_v1_to_v2( @pytest.mark.parametrize( ( "config_entry_disabled_by", + "device_disabled_by", + "entity_disabled_by", "merged_config_entry_disabled_by", "conversation_subentry_data", "main_config_entry", @@ -163,6 +165,8 @@ async def test_migration_from_v1_to_v2( [ ( [ConfigEntryDisabler.USER, None], + [DeviceEntryDisabler.CONFIG_ENTRY, None], + [RegistryEntryDisabler.CONFIG_ENTRY, None], None, [ { @@ -182,18 +186,20 @@ async def test_migration_from_v1_to_v2( ), ( [None, ConfigEntryDisabler.USER], + [None, DeviceEntryDisabler.CONFIG_ENTRY], + [None, RegistryEntryDisabler.CONFIG_ENTRY], None, [ { "conversation_entity_id": "conversation.claude", - "device_disabled_by": DeviceEntryDisabler.USER, - "entity_disabled_by": RegistryEntryDisabler.DEVICE, + "device_disabled_by": None, + "entity_disabled_by": None, "device": 0, }, { "conversation_entity_id": "conversation.claude_2", - "device_disabled_by": None, - "entity_disabled_by": None, + "device_disabled_by": DeviceEntryDisabler.USER, + "entity_disabled_by": RegistryEntryDisabler.DEVICE, "device": 1, }, ], @@ -201,6 +207,8 @@ async def test_migration_from_v1_to_v2( ), ( [ConfigEntryDisabler.USER, ConfigEntryDisabler.USER], + [DeviceEntryDisabler.CONFIG_ENTRY, DeviceEntryDisabler.CONFIG_ENTRY], + [RegistryEntryDisabler.CONFIG_ENTRY, RegistryEntryDisabler.CONFIG_ENTRY], ConfigEntryDisabler.USER, [ { @@ -211,8 +219,8 @@ async def test_migration_from_v1_to_v2( }, { "conversation_entity_id": "conversation.claude_2", - "device_disabled_by": None, - "entity_disabled_by": None, + "device_disabled_by": DeviceEntryDisabler.CONFIG_ENTRY, + "entity_disabled_by": RegistryEntryDisabler.CONFIG_ENTRY, "device": 1, }, ], @@ -225,6 +233,8 @@ async def test_migration_from_v1_disabled( device_registry: dr.DeviceRegistry, entity_registry: er.EntityRegistry, config_entry_disabled_by: list[ConfigEntryDisabler | None], + device_disabled_by: list[DeviceEntryDisabler | None], + entity_disabled_by: list[RegistryEntryDisabler | None], merged_config_entry_disabled_by: ConfigEntryDisabler | None, conversation_subentry_data: list[dict[str, Any]], main_config_entry: int, @@ -264,7 +274,7 @@ async def test_migration_from_v1_disabled( manufacturer="Anthropic", model="Claude", entry_type=dr.DeviceEntryType.SERVICE, - disabled_by=DeviceEntryDisabler.CONFIG_ENTRY, + disabled_by=device_disabled_by[0], ) entity_registry.async_get_or_create( "conversation", @@ -273,7 +283,7 @@ async def test_migration_from_v1_disabled( config_entry=mock_config_entry, device_id=device_1.id, suggested_object_id="claude", - disabled_by=RegistryEntryDisabler.CONFIG_ENTRY, + disabled_by=entity_disabled_by[0], ) device_2 = device_registry.async_get_or_create( @@ -283,6 +293,7 @@ async def test_migration_from_v1_disabled( manufacturer="Anthropic", model="Claude", entry_type=dr.DeviceEntryType.SERVICE, + disabled_by=device_disabled_by[1], ) entity_registry.async_get_or_create( "conversation", @@ -291,6 +302,7 @@ async def test_migration_from_v1_disabled( config_entry=mock_config_entry_2, device_id=device_2.id, suggested_object_id="claude", + disabled_by=entity_disabled_by[1], ) devices = [device_1, device_2] diff --git a/tests/components/google_generative_ai_conversation/test_init.py b/tests/components/google_generative_ai_conversation/test_init.py index fbd52dc9245..8098eed7f15 100644 --- a/tests/components/google_generative_ai_conversation/test_init.py +++ b/tests/components/google_generative_ai_conversation/test_init.py @@ -576,6 +576,8 @@ async def test_migration_from_v1( @pytest.mark.parametrize( ( "config_entry_disabled_by", + "device_disabled_by", + "entity_disabled_by", "merged_config_entry_disabled_by", "conversation_subentry_data", "main_config_entry", @@ -583,6 +585,8 @@ async def test_migration_from_v1( [ ( [ConfigEntryDisabler.USER, None], + [DeviceEntryDisabler.CONFIG_ENTRY, None], + [RegistryEntryDisabler.CONFIG_ENTRY, None], None, [ { @@ -602,18 +606,20 @@ async def test_migration_from_v1( ), ( [None, ConfigEntryDisabler.USER], + [None, DeviceEntryDisabler.CONFIG_ENTRY], + [None, RegistryEntryDisabler.CONFIG_ENTRY], None, [ { "conversation_entity_id": "conversation.google_generative_ai_conversation", - "device_disabled_by": DeviceEntryDisabler.USER, - "entity_disabled_by": RegistryEntryDisabler.DEVICE, + "device_disabled_by": None, + "entity_disabled_by": None, "device": 0, }, { "conversation_entity_id": "conversation.google_generative_ai_conversation_2", - "device_disabled_by": None, - "entity_disabled_by": None, + "device_disabled_by": DeviceEntryDisabler.USER, + "entity_disabled_by": RegistryEntryDisabler.DEVICE, "device": 1, }, ], @@ -621,6 +627,8 @@ async def test_migration_from_v1( ), ( [ConfigEntryDisabler.USER, ConfigEntryDisabler.USER], + [DeviceEntryDisabler.CONFIG_ENTRY, DeviceEntryDisabler.CONFIG_ENTRY], + [RegistryEntryDisabler.CONFIG_ENTRY, RegistryEntryDisabler.CONFIG_ENTRY], ConfigEntryDisabler.USER, [ { @@ -631,8 +639,8 @@ async def test_migration_from_v1( }, { "conversation_entity_id": "conversation.google_generative_ai_conversation_2", - "device_disabled_by": None, - "entity_disabled_by": None, + "device_disabled_by": DeviceEntryDisabler.CONFIG_ENTRY, + "entity_disabled_by": RegistryEntryDisabler.CONFIG_ENTRY, "device": 1, }, ], @@ -645,6 +653,8 @@ async def test_migration_from_v1_disabled( device_registry: dr.DeviceRegistry, entity_registry: er.EntityRegistry, config_entry_disabled_by: list[ConfigEntryDisabler | None], + device_disabled_by: list[DeviceEntryDisabler | None], + entity_disabled_by: list[RegistryEntryDisabler | None], merged_config_entry_disabled_by: ConfigEntryDisabler | None, conversation_subentry_data: list[dict[str, Any]], main_config_entry: int, @@ -684,7 +694,7 @@ async def test_migration_from_v1_disabled( manufacturer="Google", model="Generative AI", entry_type=dr.DeviceEntryType.SERVICE, - disabled_by=DeviceEntryDisabler.CONFIG_ENTRY, + disabled_by=device_disabled_by[0], ) entity_registry.async_get_or_create( "conversation", @@ -693,7 +703,7 @@ async def test_migration_from_v1_disabled( config_entry=mock_config_entry, device_id=device_1.id, suggested_object_id="google_generative_ai_conversation", - disabled_by=RegistryEntryDisabler.CONFIG_ENTRY, + disabled_by=entity_disabled_by[0], ) device_2 = device_registry.async_get_or_create( @@ -703,6 +713,7 @@ async def test_migration_from_v1_disabled( manufacturer="Google", model="Generative AI", entry_type=dr.DeviceEntryType.SERVICE, + disabled_by=device_disabled_by[1], ) entity_registry.async_get_or_create( "conversation", @@ -711,6 +722,7 @@ async def test_migration_from_v1_disabled( config_entry=mock_config_entry_2, device_id=device_2.id, suggested_object_id="google_generative_ai_conversation_2", + disabled_by=entity_disabled_by[1], ) devices = [device_1, device_2] diff --git a/tests/components/ollama/test_init.py b/tests/components/ollama/test_init.py index 766de8a7d6d..25e41daf276 100644 --- a/tests/components/ollama/test_init.py +++ b/tests/components/ollama/test_init.py @@ -372,6 +372,8 @@ async def test_migration_from_v1_with_same_urls( @pytest.mark.parametrize( ( "config_entry_disabled_by", + "device_disabled_by", + "entity_disabled_by", "merged_config_entry_disabled_by", "conversation_subentry_data", "main_config_entry", @@ -379,6 +381,8 @@ async def test_migration_from_v1_with_same_urls( [ ( [ConfigEntryDisabler.USER, None], + [DeviceEntryDisabler.CONFIG_ENTRY, None], + [RegistryEntryDisabler.CONFIG_ENTRY, None], None, [ { @@ -398,18 +402,20 @@ async def test_migration_from_v1_with_same_urls( ), ( [None, ConfigEntryDisabler.USER], + [None, DeviceEntryDisabler.CONFIG_ENTRY], + [None, RegistryEntryDisabler.CONFIG_ENTRY], None, [ { "conversation_entity_id": "conversation.ollama", - "device_disabled_by": DeviceEntryDisabler.USER, - "entity_disabled_by": RegistryEntryDisabler.DEVICE, + "device_disabled_by": None, + "entity_disabled_by": None, "device": 0, }, { "conversation_entity_id": "conversation.ollama_2", - "device_disabled_by": None, - "entity_disabled_by": None, + "device_disabled_by": DeviceEntryDisabler.USER, + "entity_disabled_by": RegistryEntryDisabler.DEVICE, "device": 1, }, ], @@ -417,6 +423,8 @@ async def test_migration_from_v1_with_same_urls( ), ( [ConfigEntryDisabler.USER, ConfigEntryDisabler.USER], + [DeviceEntryDisabler.CONFIG_ENTRY, DeviceEntryDisabler.CONFIG_ENTRY], + [RegistryEntryDisabler.CONFIG_ENTRY, RegistryEntryDisabler.CONFIG_ENTRY], ConfigEntryDisabler.USER, [ { @@ -427,8 +435,8 @@ async def test_migration_from_v1_with_same_urls( }, { "conversation_entity_id": "conversation.ollama_2", - "device_disabled_by": None, - "entity_disabled_by": None, + "device_disabled_by": DeviceEntryDisabler.CONFIG_ENTRY, + "entity_disabled_by": RegistryEntryDisabler.CONFIG_ENTRY, "device": 1, }, ], @@ -441,6 +449,8 @@ async def test_migration_from_v1_disabled( device_registry: dr.DeviceRegistry, entity_registry: er.EntityRegistry, config_entry_disabled_by: list[ConfigEntryDisabler | None], + device_disabled_by: list[DeviceEntryDisabler | None], + entity_disabled_by: list[RegistryEntryDisabler | None], merged_config_entry_disabled_by: ConfigEntryDisabler | None, conversation_subentry_data: list[dict[str, Any]], main_config_entry: int, @@ -474,7 +484,7 @@ async def test_migration_from_v1_disabled( manufacturer="Ollama", model="Ollama", entry_type=dr.DeviceEntryType.SERVICE, - disabled_by=DeviceEntryDisabler.CONFIG_ENTRY, + disabled_by=device_disabled_by[0], ) entity_registry.async_get_or_create( "conversation", @@ -483,7 +493,7 @@ async def test_migration_from_v1_disabled( config_entry=mock_config_entry, device_id=device_1.id, suggested_object_id="ollama", - disabled_by=RegistryEntryDisabler.CONFIG_ENTRY, + disabled_by=entity_disabled_by[0], ) device_2 = device_registry.async_get_or_create( @@ -493,6 +503,7 @@ async def test_migration_from_v1_disabled( manufacturer="Ollama", model="Ollama", entry_type=dr.DeviceEntryType.SERVICE, + disabled_by=device_disabled_by[1], ) entity_registry.async_get_or_create( "conversation", @@ -501,6 +512,7 @@ async def test_migration_from_v1_disabled( config_entry=mock_config_entry_2, device_id=device_2.id, suggested_object_id="ollama_2", + disabled_by=entity_disabled_by[1], ) devices = [device_1, device_2] diff --git a/tests/components/openai_conversation/test_init.py b/tests/components/openai_conversation/test_init.py index 66afc41826b..70d873752ae 100644 --- a/tests/components/openai_conversation/test_init.py +++ b/tests/components/openai_conversation/test_init.py @@ -868,6 +868,8 @@ async def test_migration_from_v1_with_same_keys( @pytest.mark.parametrize( ( "config_entry_disabled_by", + "device_disabled_by", + "entity_disabled_by", "merged_config_entry_disabled_by", "conversation_subentry_data", "main_config_entry", @@ -875,6 +877,8 @@ async def test_migration_from_v1_with_same_keys( [ ( [ConfigEntryDisabler.USER, None], + [DeviceEntryDisabler.CONFIG_ENTRY, None], + [RegistryEntryDisabler.CONFIG_ENTRY, None], None, [ { @@ -894,18 +898,20 @@ async def test_migration_from_v1_with_same_keys( ), ( [None, ConfigEntryDisabler.USER], + [None, DeviceEntryDisabler.CONFIG_ENTRY], + [None, RegistryEntryDisabler.CONFIG_ENTRY], None, [ { "conversation_entity_id": "conversation.chatgpt", - "device_disabled_by": DeviceEntryDisabler.USER, - "entity_disabled_by": RegistryEntryDisabler.DEVICE, + "device_disabled_by": None, + "entity_disabled_by": None, "device": 0, }, { "conversation_entity_id": "conversation.chatgpt_2", - "device_disabled_by": None, - "entity_disabled_by": None, + "device_disabled_by": DeviceEntryDisabler.USER, + "entity_disabled_by": RegistryEntryDisabler.DEVICE, "device": 1, }, ], @@ -913,6 +919,8 @@ async def test_migration_from_v1_with_same_keys( ), ( [ConfigEntryDisabler.USER, ConfigEntryDisabler.USER], + [DeviceEntryDisabler.CONFIG_ENTRY, DeviceEntryDisabler.CONFIG_ENTRY], + [RegistryEntryDisabler.CONFIG_ENTRY, RegistryEntryDisabler.CONFIG_ENTRY], ConfigEntryDisabler.USER, [ { @@ -923,8 +931,8 @@ async def test_migration_from_v1_with_same_keys( }, { "conversation_entity_id": "conversation.chatgpt_2", - "device_disabled_by": None, - "entity_disabled_by": None, + "device_disabled_by": DeviceEntryDisabler.CONFIG_ENTRY, + "entity_disabled_by": RegistryEntryDisabler.CONFIG_ENTRY, "device": 1, }, ], @@ -937,6 +945,8 @@ async def test_migration_from_v1_disabled( device_registry: dr.DeviceRegistry, entity_registry: er.EntityRegistry, config_entry_disabled_by: list[ConfigEntryDisabler | None], + device_disabled_by: list[DeviceEntryDisabler | None], + entity_disabled_by: list[RegistryEntryDisabler | None], merged_config_entry_disabled_by: ConfigEntryDisabler | None, conversation_subentry_data: list[dict[str, Any]], main_config_entry: int, @@ -976,7 +986,7 @@ async def test_migration_from_v1_disabled( manufacturer="OpenAI", model="ChatGPT", entry_type=dr.DeviceEntryType.SERVICE, - disabled_by=DeviceEntryDisabler.CONFIG_ENTRY, + disabled_by=device_disabled_by[0], ) entity_registry.async_get_or_create( "conversation", @@ -985,7 +995,7 @@ async def test_migration_from_v1_disabled( config_entry=mock_config_entry, device_id=device_1.id, suggested_object_id="chatgpt", - disabled_by=RegistryEntryDisabler.CONFIG_ENTRY, + disabled_by=entity_disabled_by[0], ) device_2 = device_registry.async_get_or_create( @@ -995,6 +1005,7 @@ async def test_migration_from_v1_disabled( manufacturer="OpenAI", model="ChatGPT", entry_type=dr.DeviceEntryType.SERVICE, + disabled_by=device_disabled_by[1], ) entity_registry.async_get_or_create( "conversation", @@ -1003,6 +1014,7 @@ async def test_migration_from_v1_disabled( config_entry=mock_config_entry_2, device_id=device_2.id, suggested_object_id="chatgpt_2", + disabled_by=entity_disabled_by[1], ) devices = [device_1, device_2] diff --git a/tests/helpers/test_device_registry.py b/tests/helpers/test_device_registry.py index 30fb22b8e09..80910d42630 100644 --- a/tests/helpers/test_device_registry.py +++ b/tests/helpers/test_device_registry.py @@ -3279,6 +3279,266 @@ async def test_update_suggested_area( assert updated_entry.area_id == device_area_id +@pytest.mark.parametrize( + ( + "new_config_entry_disabled_by", + "device_disabled_by_initial", + "device_disabled_by_updated", + "extra_changes", + ), + [ + ( + None, + None, + None, + {}, + ), + # Config entry not disabled, device was disabled by config entry. + # Device not disabled when updated. + ( + None, + dr.DeviceEntryDisabler.CONFIG_ENTRY, + None, + {"disabled_by": dr.DeviceEntryDisabler.CONFIG_ENTRY}, + ), + ( + None, + dr.DeviceEntryDisabler.INTEGRATION, + dr.DeviceEntryDisabler.INTEGRATION, + {}, + ), + ( + None, + dr.DeviceEntryDisabler.USER, + dr.DeviceEntryDisabler.USER, + {}, + ), + ( + config_entries.ConfigEntryDisabler.USER, + None, + None, + {}, + ), + ( + config_entries.ConfigEntryDisabler.USER, + dr.DeviceEntryDisabler.CONFIG_ENTRY, + dr.DeviceEntryDisabler.CONFIG_ENTRY, + {}, + ), + ( + config_entries.ConfigEntryDisabler.USER, + dr.DeviceEntryDisabler.INTEGRATION, + dr.DeviceEntryDisabler.INTEGRATION, + {}, + ), + ( + config_entries.ConfigEntryDisabler.USER, + dr.DeviceEntryDisabler.USER, + dr.DeviceEntryDisabler.USER, + {}, + ), + ], +) +@pytest.mark.usefixtures("freezer") +async def test_update_add_config_entry_disabled_by( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + new_config_entry_disabled_by: config_entries.ConfigEntryDisabler | None, + device_disabled_by_initial: dr.DeviceEntryDisabler | None, + device_disabled_by_updated: dr.DeviceEntryDisabler | None, + extra_changes: dict[str, Any], +) -> None: + """Check how the disabled_by flag is treated when adding a config entry.""" + config_entry_1 = MockConfigEntry(title=None) + config_entry_1.add_to_hass(hass) + config_entry_2 = MockConfigEntry( + title=None, disabled_by=new_config_entry_disabled_by + ) + config_entry_2.add_to_hass(hass) + update_events = async_capture_events(hass, dr.EVENT_DEVICE_REGISTRY_UPDATED) + entry = device_registry.async_get_or_create( + config_entry_id=config_entry_1.entry_id, + config_subentry_id=None, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + disabled_by=device_disabled_by_initial, + ) + assert entry.disabled_by == device_disabled_by_initial + + entry2 = device_registry.async_update_device( + entry.id, add_config_entry_id=config_entry_2.entry_id + ) + + assert entry2 == dr.DeviceEntry( + config_entries={config_entry_1.entry_id, config_entry_2.entry_id}, + config_entries_subentries={ + config_entry_1.entry_id: {None}, + config_entry_2.entry_id: {None}, + }, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:ab:cd:ef")}, + created_at=utcnow(), + disabled_by=device_disabled_by_updated, + id=entry.id, + modified_at=utcnow(), + primary_config_entry=None, + ) + + await hass.async_block_till_done() + + assert len(update_events) == 2 + assert update_events[0].data == { + "action": "create", + "device_id": entry.id, + } + assert update_events[1].data == { + "action": "update", + "device_id": entry.id, + "changes": { + "config_entries": {config_entry_1.entry_id}, + "config_entries_subentries": {config_entry_1.entry_id: {None}}, + } + | extra_changes, + } + + +@pytest.mark.parametrize( + ( + "removed_config_entry_disabled_by", + "device_disabled_by_initial", + "device_disabled_by_updated", + "extra_changes", + ), + [ + # The non-disabled config entry is removed, device changed to + # disabled by config entry. + ( + None, + None, + dr.DeviceEntryDisabler.CONFIG_ENTRY, + {"disabled_by": None}, + ), + ( + None, + dr.DeviceEntryDisabler.CONFIG_ENTRY, + dr.DeviceEntryDisabler.CONFIG_ENTRY, + {}, + ), + ( + None, + dr.DeviceEntryDisabler.INTEGRATION, + dr.DeviceEntryDisabler.INTEGRATION, + {}, + ), + ( + None, + dr.DeviceEntryDisabler.USER, + dr.DeviceEntryDisabler.USER, + {}, + ), + # In this test, the device is in an invalid state: config entry disabled, + # device not disabled. After removing the config entry, the device is disabled + # by checking the remaining config entry. + ( + config_entries.ConfigEntryDisabler.USER, + None, + dr.DeviceEntryDisabler.CONFIG_ENTRY, + {"disabled_by": None}, + ), + ( + config_entries.ConfigEntryDisabler.USER, + dr.DeviceEntryDisabler.CONFIG_ENTRY, + dr.DeviceEntryDisabler.CONFIG_ENTRY, + {}, + ), + ( + config_entries.ConfigEntryDisabler.USER, + dr.DeviceEntryDisabler.INTEGRATION, + dr.DeviceEntryDisabler.INTEGRATION, + {}, + ), + ( + config_entries.ConfigEntryDisabler.USER, + dr.DeviceEntryDisabler.USER, + dr.DeviceEntryDisabler.USER, + {}, + ), + ], +) +@pytest.mark.usefixtures("freezer") +async def test_update_remove_config_entry_disabled_by( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + removed_config_entry_disabled_by: config_entries.ConfigEntryDisabler | None, + device_disabled_by_initial: dr.DeviceEntryDisabler | None, + device_disabled_by_updated: dr.DeviceEntryDisabler | None, + extra_changes: dict[str, Any], +) -> None: + """Check how the disabled_by flag is treated when removing a config entry.""" + config_entry_1 = MockConfigEntry( + title=None, disabled_by=removed_config_entry_disabled_by + ) + config_entry_1.add_to_hass(hass) + config_entry_2 = MockConfigEntry( + title=None, disabled_by=config_entries.ConfigEntryDisabler.USER + ) + config_entry_2.add_to_hass(hass) + update_events = async_capture_events(hass, dr.EVENT_DEVICE_REGISTRY_UPDATED) + entry = device_registry.async_get_or_create( + config_entry_id=config_entry_1.entry_id, + config_subentry_id=None, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + disabled_by=device_disabled_by_initial, + ) + assert entry.disabled_by == device_disabled_by_initial + + entry2 = device_registry.async_update_device( + entry.id, add_config_entry_id=config_entry_2.entry_id + ) + assert entry2.disabled_by == device_disabled_by_initial + + entry3 = device_registry.async_update_device( + entry.id, remove_config_entry_id=config_entry_1.entry_id + ) + + assert entry3 == dr.DeviceEntry( + config_entries={config_entry_2.entry_id}, + config_entries_subentries={config_entry_2.entry_id: {None}}, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:ab:cd:ef")}, + created_at=utcnow(), + disabled_by=device_disabled_by_updated, + id=entry.id, + modified_at=utcnow(), + primary_config_entry=None, + ) + + await hass.async_block_till_done() + + assert len(update_events) == 3 + assert update_events[0].data == { + "action": "create", + "device_id": entry.id, + } + assert update_events[1].data == { + "action": "update", + "device_id": entry.id, + "changes": { + "config_entries": {config_entry_1.entry_id}, + "config_entries_subentries": {config_entry_1.entry_id: {None}}, + }, + } + assert update_events[2].data == { + "action": "update", + "device_id": entry.id, + "changes": { + "config_entries": {config_entry_1.entry_id, config_entry_2.entry_id}, + "config_entries_subentries": { + config_entry_1.entry_id: {None}, + config_entry_2.entry_id: {None}, + }, + } + | extra_changes, + } + + async def test_cleanup_device_registry( hass: HomeAssistant, device_registry: dr.DeviceRegistry,