Add updates argument to update_reload_and_abort helper (#127781)

* Add updates argument to update_reload_and_abort helper

* Also apply to airvisual_pro

* Rename argument

* docstring

* Use modern syntax

Co-authored-by: Erik Montnemery <erik@montnemery.com>

* Apply suggestion

Co-authored-by: Erik Montnemery <erik@montnemery.com>

* Apply suggestion

* Docstring

---------

Co-authored-by: Erik Montnemery <erik@montnemery.com>

homeassistant/components/airvisual_pro/config_flow.py

@@ -14,7 +14,7 @@ from pyairvisual.node import (
 )
 import voluptuous as vol
 
-from homeassistant.config_entries import ConfigEntry, ConfigFlow, ConfigFlowResult
+from homeassistant.config_entries import ConfigFlow, ConfigFlowResult
 from homeassistant.const import CONF_IP_ADDRESS, CONF_PASSWORD
 
 from .const import DOMAIN, LOGGER
@@ -76,7 +76,7 @@ class AirVisualProFlowHandler(ConfigFlow, domain=DOMAIN):
 
     VERSION = 1
 
-    _reauth_entry: ConfigEntry
+    _reauth_entry_data: Mapping[str, Any]
 
     async def async_step_import(self, import_data: dict[str, Any]) -> ConfigFlowResult:
         """Import a config entry from `airvisual` integration (see #83882)."""
@@ -86,7 +86,7 @@ class AirVisualProFlowHandler(ConfigFlow, domain=DOMAIN):
         self, entry_data: Mapping[str, Any]
     ) -> ConfigFlowResult:
         """Handle configuration by re-auth."""
-        self._reauth_entry = self._get_reauth_entry()
+        self._reauth_entry_data = entry_data
         return await self.async_step_reauth_confirm()
 
     async def async_step_reauth_confirm(
@@ -99,7 +99,7 @@ class AirVisualProFlowHandler(ConfigFlow, domain=DOMAIN):
             )
 
         validation_result = await async_validate_credentials(
-            self._reauth_entry.data[CONF_IP_ADDRESS], user_input[CONF_PASSWORD]
+            self._reauth_entry_data[CONF_IP_ADDRESS], user_input[CONF_PASSWORD]
         )
 
         if validation_result.errors:
@@ -110,7 +110,7 @@ class AirVisualProFlowHandler(ConfigFlow, domain=DOMAIN):
             )
 
         return self.async_update_reload_and_abort(
-            self._reauth_entry, data=self._reauth_entry.data | user_input
+            self._get_reauth_entry(), data_updates=user_input
         )
 
     async def async_step_user(

homeassistant/components/aosmith/config_flow.py

@@ -88,12 +88,11 @@ class AOSmithConfigFlow(ConfigFlow, domain=DOMAIN):
         if user_input:
             password = user_input[CONF_PASSWORD]
 
-            entry = self._get_reauth_entry()
             error = await self._async_validate_credentials(self._reauth_email, password)
             if error is None:
                 return self.async_update_reload_and_abort(
-                    entry,
+                    self._get_reauth_entry(),
-                    data=entry.data | user_input,
+                    data_updates=user_input,
                 )
             errors["base"] = error
 

homeassistant/config_entries.py

@@ -2761,11 +2761,30 @@ class ConfigFlow(ConfigEntryBaseFlow):
         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,
         reload_even_if_entry_is_unchanged: bool = True,
     ) -> ConfigFlowResult:
-        """Update config entry, reload config entry and finish config flow."""
+        """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
+        """
+        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(
             entry=entry,
             unique_id=unique_id,

tests/test_config_entries.py

@@ -5122,6 +5122,7 @@ def test_raise_trying_to_add_same_config_entry_twice(
         "expected_data",
         "expected_options",
         "calls_entry_load_unload",
+        "raises",
     ),
     [
         (
@@ -5136,6 +5137,7 @@ def test_raise_trying_to_add_same_config_entry_twice(
             {"vendor": "data2"},
             {"vendor": "options2"},
             (2, 1),
+            None,
         ),
         (
             {
@@ -5149,6 +5151,7 @@ def test_raise_trying_to_add_same_config_entry_twice(
             {"vendor": "data"},
             {"vendor": "options"},
             (2, 1),
+            None,
         ),
         (
             {
@@ -5163,6 +5166,7 @@ def test_raise_trying_to_add_same_config_entry_twice(
             {"vendor": "data2"},
             {"vendor": "options2"},
             (2, 1),
+            None,
         ),
         (
             {
@@ -5177,6 +5181,7 @@ def test_raise_trying_to_add_same_config_entry_twice(
             {"vendor": "data"},
             {"vendor": "options"},
             (1, 0),
+            None,
         ),
         (
             {},
@@ -5185,6 +5190,7 @@ def test_raise_trying_to_add_same_config_entry_twice(
             {"vendor": "data"},
             {"vendor": "options"},
             (2, 1),
+            None,
         ),
         (
             {"data": {"buyer": "me"}, "options": {}},
@@ -5193,6 +5199,31 @@ def test_raise_trying_to_add_same_config_entry_twice(
             {"buyer": "me"},
             {},
             (2, 1),
+            None,
+        ),
+        (
+            {"data_updates": {"buyer": "me"}},
+            "Test",
+            "1234",
+            {"vendor": "data", "buyer": "me"},
+            {"vendor": "options"},
+            (2, 1),
+            None,
+        ),
+        (
+            {
+                "unique_id": "5678",
+                "title": "Updated title",
+                "data": {"vendor": "data2"},
+                "options": {"vendor": "options2"},
+                "data_updates": {"buyer": "me"},
+            },
+            "Test",
+            "1234",
+            {"vendor": "data"},
+            {"vendor": "options"},
+            (1, 0),
+            ValueError,
         ),
     ],
     ids=[
@@ -5202,6 +5233,8 @@ def test_raise_trying_to_add_same_config_entry_twice(
         "unchanged_entry_no_reload",
         "no_kwargs",
         "replace_data",
+        "update_data",
+        "update_and_data_raises",
     ],
 )
 @pytest.mark.parametrize(
@@ -5221,6 +5254,7 @@ async def test_update_entry_and_reload(
     expected_options: dict[str, Any],
     kwargs: dict[str, Any],
     calls_entry_load_unload: tuple[int, int],
+    raises: type[Exception] | None,
 ) -> None:
     """Test updating an entry and reloading."""
     entry = MockConfigEntry(
@@ -5255,11 +5289,15 @@ async def test_update_entry_and_reload(
             """Mock Reconfigure."""
             return self.async_update_reload_and_abort(entry, **kwargs)
 
+    err: Exception
     with mock_config_flow("comp", MockFlowHandler):
+        try:
             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)
+        except Exception as ex:  # noqa: BLE001
+            err = ex
 
     await hass.async_block_till_done()
 
@@ -5268,6 +5306,9 @@ async def test_update_entry_and_reload(
     assert entry.data == expected_data
     assert entry.options == expected_options
     assert entry.state == config_entries.ConfigEntryState.LOADED
+    if raises:
+        assert isinstance(err, raises)
+    else:
         assert result["type"] == FlowResultType.ABORT
         assert result["reason"] == reason
     # Assert entry was reloaded