Add async_update_and_abort method to config flow (home-assistant#153146)

Co-authored-by: Erik Montnemery <[email protected]>

Author: gjohansson-ST

homeassistant/config_entries.py

@@ -3210,6 +3210,76 @@ def async_create_entry(  # type: ignore[override]
 
         return result
 
+    @callback
+    def __async_update(
+        self,
+        entry: ConfigEntry,
+        *,
+        unique_id: str | None | UndefinedType,
+        title: str | UndefinedType,
+        data: Mapping[str, Any] | UndefinedType,
+        data_updates: Mapping[str, Any] | UndefinedType,
+        options: Mapping[str, Any] | UndefinedType,
+    ) -> bool:
+        """Update config entry and return result.
+
+        Internal to be used by update_and_abort and update_reload_and_abort methods only.
+        """
+
+        if data_updates is not UNDEFINED:
+            if data is not UNDEFINED:
+                raise ValueError("Cannot set both data and data_updates")
+            data = entry.data | data_updates
+        return self.hass.config_entries.async_update_entry(
+            entry=entry,
+            unique_id=unique_id,
+            title=title,
+            data=data,
+            options=options,
+        )
+
+    @callback
+    def async_update_and_abort(
+        self,
+        entry: ConfigEntry,
+        *,
+        unique_id: str | None | UndefinedType = UNDEFINED,
+        title: str | UndefinedType = UNDEFINED,
+        data: Mapping[str, Any] | UndefinedType = UNDEFINED,
+        data_updates: Mapping[str, Any] | UndefinedType = UNDEFINED,
+        options: Mapping[str, Any] | UndefinedType = UNDEFINED,
+        reason: str | UndefinedType = UNDEFINED,
+    ) -> ConfigFlowResult:
+        """Update config entry and finish config flow.
+
+        Args:
+            entry: config entry to update
+            unique_id: replace the unique_id of the entry
+            title: replace the title of the entry
+            data: replace the entry data with new data
+            data_updates: add items from data_updates to entry data - existing keys
+                are overridden
+            options: replace the entry options with new options
+            reason: set the reason for the abort, defaults to
+                `reauth_successful` or `reconfigure_successful` based on flow source
+
+        Returns:
+            ConfigFlowResult: The result of the config flow.
+        """
+        self.__async_update(
+            entry=entry,
+            unique_id=unique_id,
+            title=title,
+            data=data,
+            data_updates=data_updates,
+            options=options,
+        )
+        if reason is UNDEFINED:
+            reason = "reauth_successful"
+            if self.source == SOURCE_RECONFIGURE:
+                reason = "reconfigure_successful"
+        return self.async_abort(reason=reason)
+
     @callback
     def async_update_reload_and_abort(
         self,
@@ -3225,28 +3295,28 @@ def async_update_reload_and_abort(
     ) -> ConfigFlowResult:
         """Update config entry, reload config entry and finish config flow.
 
-        :param data: replace the entry data with new data
-        :param data_updates: add items from data_updates to entry data - existing keys
-        are overridden
-        :param options: replace the entry options with new options
-        :param title: replace the title of the entry
-        :param unique_id: replace the unique_id of the entry
-
-        :param reason: set the reason for the abort, defaults to
-        `reauth_successful` or `reconfigure_successful` based on flow source
-
-        :param reload_even_if_entry_is_unchanged: set this to `False` if the entry
-        should not be reloaded if it is unchanged
+        Args:
+            entry: config entry to update and reload
+            unique_id: replace the unique_id of the entry
+            title: replace the title of the entry
+            data: replace the entry data with new data
+            data_updates: add items from data_updates to entry data - existing keys
+                are overridden
+            options: replace the entry options with new options
+            reason: set the reason for the abort, defaults to
+                `reauth_successful` or `reconfigure_successful` based on flow source
+            reload_even_if_entry_is_unchanged: set this to `False` if the entry
+                should not be reloaded if it is unchanged
+
+        Returns:
+            ConfigFlowResult: The result of the config flow.
         """
-        if data_updates is not UNDEFINED:
-            if data is not UNDEFINED:
-                raise ValueError("Cannot set both data and data_updates")
-            data = entry.data | data_updates
-        result = self.hass.config_entries.async_update_entry(
+        result = self.__async_update(
             entry=entry,
             unique_id=unique_id,
             title=title,
             data=data,
+            data_updates=data_updates,
             options=options,
         )
         if reload_even_if_entry_is_unchanged or result:

tests/test_config_entries.py

@@ -6626,6 +6626,83 @@ async def async_step_reconfigure(self, data):
     assert len(comp.async_unload_entry.mock_calls) == calls_entry_load_unload[1]
 
 
+@pytest.mark.parametrize(
+    ("source", "reason"),
+    [
+        (config_entries.SOURCE_REAUTH, "reauth_successful"),
+        (config_entries.SOURCE_RECONFIGURE, "reconfigure_successful"),
+    ],
+)
+async def test_update_entry_without_reload(
+    hass: HomeAssistant,
+    source: str,
+    reason: str,
+) -> None:
+    """Test updating an entry without reloading."""
+    entry = MockConfigEntry(
+        domain="comp",
+        unique_id="1234",
+        title="Test",
+        data={"vendor": "data"},
+        options={"vendor": "options"},
+    )
+    entry.add_to_hass(hass)
+
+    comp = MockModule(
+        "comp",
+        async_setup_entry=AsyncMock(return_value=True),
+        async_unload_entry=AsyncMock(return_value=True),
+    )
+    mock_integration(hass, comp)
+    mock_platform(hass, "comp.config_flow", None)
+
+    await hass.config_entries.async_setup(entry.entry_id)
+
+    class MockFlowHandler(config_entries.ConfigFlow):
+        """Define a mock flow handler."""
+
+        VERSION = 1
+
+        async def async_step_reauth(self, data):
+            """Mock Reauth."""
+            return self.async_update_and_abort(
+                entry,
+                unique_id="5678",
+                title="Updated title",
+                data={"vendor": "data2"},
+                options={"vendor": "options2"},
+            )
+
+        async def async_step_reconfigure(self, data):
+            """Mock Reconfigure."""
+            return self.async_update_and_abort(
+                entry,
+                unique_id="5678",
+                title="Updated title",
+                data={"vendor": "data2"},
+                options={"vendor": "options2"},
+            )
+
+    with mock_config_flow("comp", MockFlowHandler):
+        if source == config_entries.SOURCE_REAUTH:
+            result = await entry.start_reauth_flow(hass)
+        elif source == config_entries.SOURCE_RECONFIGURE:
+            result = await entry.start_reconfigure_flow(hass)
+
+    await hass.async_block_till_done()
+
+    assert entry.title == "Updated title"
+    assert entry.unique_id == "5678"
+    assert entry.data == {"vendor": "data2"}
+    assert entry.options == {"vendor": "options2"}
+    assert entry.state == config_entries.ConfigEntryState.LOADED
+    assert result["type"] == FlowResultType.ABORT
+    assert result["reason"] == reason
+    # Assert entry is not reloaded
+    assert len(comp.async_setup_entry.mock_calls) == 1
+    assert len(comp.async_unload_entry.mock_calls) == 0
+
+
 @pytest.mark.parametrize(
     (
         "kwargs",