From 21266e1c685dc1c668acd03803f5573198f1fea5 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Wed, 2 Oct 2024 12:00:07 +0200 Subject: [PATCH] Add config_flow helper to get reauth/reconfigure config entry (#127115) * Add config_flow helper to get config entry from context * Simplify * Apply to aussie_broadband * Another example * Rename and adjust docstring * Simplify * Add test * Refactor to hide context * Raise * Improve coverage * Use AttributeError * Use ValueError * Raise UnknownEntry --- .../aussie_broadband/config_flow.py | 5 +- .../bryant_evolution/config_flow.py | 5 +- homeassistant/config_entries.py | 30 ++++ tests/test_config_entries.py | 146 ++++++++++++++++++ 4 files changed, 178 insertions(+), 8 deletions(-) diff --git a/homeassistant/components/aussie_broadband/config_flow.py b/homeassistant/components/aussie_broadband/config_flow.py index 65507d57e8b..540c04f3993 100644 --- a/homeassistant/components/aussie_broadband/config_flow.py +++ b/homeassistant/components/aussie_broadband/config_flow.py @@ -99,10 +99,7 @@ class AussieBroadbandConfigFlow(ConfigFlow, domain=DOMAIN): } if not (errors := await self.async_auth(data)): - entry = self.hass.config_entries.async_get_entry( - self.context["entry_id"] - ) - assert entry + entry = self._get_reauth_entry() return self.async_update_reload_and_abort(entry, data=data) return self.async_show_form( diff --git a/homeassistant/components/bryant_evolution/config_flow.py b/homeassistant/components/bryant_evolution/config_flow.py index 9cfb9b2ec7e..7d85406b707 100644 --- a/homeassistant/components/bryant_evolution/config_flow.py +++ b/homeassistant/components/bryant_evolution/config_flow.py @@ -75,10 +75,7 @@ class BryantConfigFlow(ConfigFlow, domain=DOMAIN): if user_input is not None: system_zone = await _enumerate_sz(user_input[CONF_FILENAME]) if len(system_zone) != 0: - our_entry = self.hass.config_entries.async_get_entry( - self.context["entry_id"] - ) - assert our_entry is not None, "Could not find own entry" + our_entry = self._get_reconfigure_entry() return self.async_update_reload_and_abort( entry=our_entry, data={ diff --git a/homeassistant/config_entries.py b/homeassistant/config_entries.py index f4ef5cd3ad1..906303ec95b 100644 --- a/homeassistant/config_entries.py +++ b/homeassistant/config_entries.py @@ -2726,6 +2726,36 @@ class ConfigFlow(ConfigEntryBaseFlow): """Return True if other_flow is matching this flow.""" raise NotImplementedError + @property + def _reauth_entry_id(self) -> str: + """Return reauth entry id.""" + if self.source != SOURCE_REAUTH: + raise ValueError(f"Source is {self.source}, expected {SOURCE_REAUTH}") + return self.context["entry_id"] # type: ignore[no-any-return] + + @callback + def _get_reauth_entry(self) -> ConfigEntry: + """Return the reauth config entry linked to the current context.""" + if entry := self.hass.config_entries.async_get_entry(self._reauth_entry_id): + return entry + raise UnknownEntry + + @property + def _reconfigure_entry_id(self) -> str: + """Return reconfigure entry id.""" + if self.source != SOURCE_RECONFIGURE: + raise ValueError(f"Source is {self.source}, expected {SOURCE_RECONFIGURE}") + return self.context["entry_id"] # type: ignore[no-any-return] + + @callback + def _get_reconfigure_entry(self) -> ConfigEntry: + """Return the reconfigure config entry linked to the current context.""" + if entry := self.hass.config_entries.async_get_entry( + self._reconfigure_entry_id + ): + return entry + raise UnknownEntry + class OptionsFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]): """Flow to set options for a configuration entry.""" diff --git a/tests/test_config_entries.py b/tests/test_config_entries.py index 59ea7cdd808..b7666516644 100644 --- a/tests/test_config_entries.py +++ b/tests/test_config_entries.py @@ -6383,6 +6383,152 @@ async def test_async_has_matching_flow_not_implemented( manager.flow.async_has_matching_flow(flow) +async def test_get_reauth_entry( + hass: HomeAssistant, manager: config_entries.ConfigEntries +) -> None: + """Test _get_context_entry behavior.""" + entry = MockConfigEntry( + title="test_title", + domain="test", + entry_id="01J915Q6T9F6G5V0QJX6HBC94T", + data={"host": "any", "port": 123}, + unique_id=None, + ) + entry.add_to_hass(hass) + + mock_integration(hass, MockModule("test")) + mock_platform(hass, "test.config_flow", None) + + class TestFlow(config_entries.ConfigFlow): + VERSION = 1 + + async def async_step_user(self, user_input=None): + """Test user step.""" + return await self._async_step_confirm() + + async def async_step_reauth(self, entry_data): + """Test reauth step.""" + return await self._async_step_confirm() + + async def async_step_reconfigure(self, entry_data): + """Test reauth step.""" + return await self._async_step_confirm() + + async def _async_step_confirm(self): + """Confirm input.""" + try: + entry = self._get_reauth_entry() + except ValueError as err: + reason = str(err) + except config_entries.UnknownEntry: + reason = "Entry not found" + else: + reason = f"Found entry {entry.title}" + try: + entry_id = self._reauth_entry_id + except ValueError: + reason = f"{reason}: -" + else: + reason = f"{reason}: {entry_id}" + return self.async_abort(reason=reason) + + # A reauth flow finds the config entry from context + with mock_config_flow("test", TestFlow): + result = await entry.start_reauth_flow(hass) + assert result["reason"] == "Found entry test_title: 01J915Q6T9F6G5V0QJX6HBC94T" + + # The config entry is removed before the reauth flow is aborted + with mock_config_flow("test", TestFlow): + result = await entry.start_reauth_flow(hass, context={"entry_id": "01JRemoved"}) + assert result["reason"] == "Entry not found: 01JRemoved" + + # A reconfigure flow does not have access to the config entry + with mock_config_flow("test", TestFlow): + result = await entry.start_reconfigure_flow(hass) + assert result["reason"] == "Source is reconfigure, expected reauth: -" + + # A user flow does not have access to the config entry + with mock_config_flow("test", TestFlow): + result = await manager.flow.async_init( + "test", context={"source": config_entries.SOURCE_USER} + ) + assert result["reason"] == "Source is user, expected reauth: -" + + +async def test_get_reconfigure_entry( + hass: HomeAssistant, manager: config_entries.ConfigEntries +) -> None: + """Test _get_context_entry behavior.""" + entry = MockConfigEntry( + title="test_title", + domain="test", + entry_id="01J915Q6T9F6G5V0QJX6HBC94T", + data={"host": "any", "port": 123}, + unique_id=None, + ) + entry.add_to_hass(hass) + + mock_integration(hass, MockModule("test")) + mock_platform(hass, "test.config_flow", None) + + class TestFlow(config_entries.ConfigFlow): + VERSION = 1 + + async def async_step_user(self, user_input=None): + """Test user step.""" + return await self._async_step_confirm() + + async def async_step_reauth(self, entry_data): + """Test reauth step.""" + return await self._async_step_confirm() + + async def async_step_reconfigure(self, entry_data): + """Test reauth step.""" + return await self._async_step_confirm() + + async def _async_step_confirm(self): + """Confirm input.""" + try: + entry = self._get_reconfigure_entry() + except ValueError as err: + reason = str(err) + except config_entries.UnknownEntry: + reason = "Entry not found" + else: + reason = f"Found entry {entry.title}" + try: + entry_id = self._reconfigure_entry_id + except ValueError: + reason = f"{reason}: -" + else: + reason = f"{reason}: {entry_id}" + return self.async_abort(reason=reason) + + # A reauth flow does not have access to the config entry from context + with mock_config_flow("test", TestFlow): + result = await entry.start_reauth_flow(hass) + assert result["reason"] == "Source is reauth, expected reconfigure: -" + + # A reconfigure flow finds the config entry + with mock_config_flow("test", TestFlow): + result = await entry.start_reconfigure_flow(hass) + assert result["reason"] == "Found entry test_title: 01J915Q6T9F6G5V0QJX6HBC94T" + + # A reconfigure flow finds the config entry + with mock_config_flow("test", TestFlow): + result = await entry.start_reconfigure_flow( + hass, context={"entry_id": "01JRemoved"} + ) + assert result["reason"] == "Entry not found: 01JRemoved" + + # A user flow does not have access to the config entry + with mock_config_flow("test", TestFlow): + result = await manager.flow.async_init( + "test", context={"source": config_entries.SOURCE_USER} + ) + assert result["reason"] == "Source is user, expected reconfigure: -" + + async def test_reauth_helper_alignment( hass: HomeAssistant, manager: config_entries.ConfigEntries,