From 663482fb10e0525ea1a5efdb49ea3f7a183a2f6d Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Wed, 30 Nov 2022 10:18:57 +0100
Subject: [PATCH] Add flow_state to SchemaCommonFlowHandler (#82967)

* Add local_context to SchemaCommonFlowHandler

* Rename to context

* Rename to flow_state
---
 .../helpers/schema_config_entry_flow.py       | 10 +++
 .../helpers/test_schema_config_entry_flow.py  | 77 +++++++++++++++++++
 2 files changed, 87 insertions(+)

diff --git a/homeassistant/helpers/schema_config_entry_flow.py b/homeassistant/helpers/schema_config_entry_flow.py
index 8d72910b6fe..a730a311db9 100644
--- a/homeassistant/helpers/schema_config_entry_flow.py
+++ b/homeassistant/helpers/schema_config_entry_flow.py
@@ -98,6 +98,7 @@ class SchemaCommonFlowHandler:
         self._flow = flow
         self._handler = handler
         self._options = options if options is not None else {}
+        self._flow_state: dict[str, Any] = {}
 
     @property
     def parent_handler(self) -> SchemaConfigFlowHandler | SchemaOptionsFlowHandler:
@@ -109,6 +110,15 @@ class SchemaCommonFlowHandler:
         """Return the options linked to the current flow handler."""
         return self._options
 
+    @property
+    def flow_state(self) -> dict[str, Any]:
+        """Return the flow state, used to store temporary data.
+
+        It can be used for example to store the key or the index of a sub-item
+        that will be edited in the next step.
+        """
+        return self._flow_state
+
     async def async_step(
         self, step_id: str, user_input: dict[str, Any] | None = None
     ) -> FlowResult:
diff --git a/tests/helpers/test_schema_config_entry_flow.py b/tests/helpers/test_schema_config_entry_flow.py
index 27f4aa0d220..29d01b2d2b6 100644
--- a/tests/helpers/test_schema_config_entry_flow.py
+++ b/tests/helpers/test_schema_config_entry_flow.py
@@ -17,6 +17,7 @@ from homeassistant.helpers.schema_config_entry_flow import (
     SchemaFlowError,
     SchemaFlowFormStep,
     SchemaFlowMenuStep,
+    SchemaOptionsFlowHandler,
     wrapped_entity_config_entry_title,
 )
 from homeassistant.util.decorator import Registry
@@ -555,3 +556,79 @@ async def test_suggested_values(
         result["flow_id"], {"option1": "blabla"}
     )
     assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY
+
+
+async def test_options_flow_state(hass: HomeAssistant) -> None:
+    """Test flow_state handling in SchemaFlowFormStep."""
+
+    OPTIONS_SCHEMA = vol.Schema(
+        {vol.Optional("option1", default="a very reasonable default"): str}
+    )
+
+    def _init_schema(handler: SchemaCommonFlowHandler) -> None:
+        handler.flow_state["idx"] = None
+
+    def _validate_step1_input(
+        handler: SchemaCommonFlowHandler, user_input: dict[str, Any]
+    ) -> dict[str, Any]:
+        handler.flow_state["idx"] = user_input["option1"]
+        return user_input
+
+    def _validate_step2_input(
+        handler: SchemaCommonFlowHandler, user_input: dict[str, Any]
+    ) -> dict[str, Any]:
+        user_input["idx_from_flow_state"] = handler.flow_state["idx"]
+        return user_input
+
+    OPTIONS_FLOW: dict[str, SchemaFlowFormStep | SchemaFlowMenuStep] = {
+        "init": SchemaFlowFormStep(_init_schema, next_step="step_1"),
+        "step_1": SchemaFlowFormStep(
+            OPTIONS_SCHEMA,
+            validate_user_input=_validate_step1_input,
+            next_step="step_2",
+        ),
+        "step_2": SchemaFlowFormStep(
+            OPTIONS_SCHEMA,
+            validate_user_input=_validate_step2_input,
+        ),
+    }
+
+    class TestFlow(SchemaConfigFlowHandler, domain="test"):
+        config_flow = {}
+        options_flow = OPTIONS_FLOW
+
+    config_entry = MockConfigEntry(
+        data={},
+        domain="test",
+        options={"option1": "initial value"},
+    )
+    config_entry.add_to_hass(hass)
+
+    # Start flow in basic mode, flow state is initialised with None value
+    result = await hass.config_entries.options.async_init(config_entry.entry_id)
+    assert result["type"] == data_entry_flow.FlowResultType.FORM
+    assert result["step_id"] == "step_1"
+
+    options_handler: SchemaOptionsFlowHandler
+    options_handler = hass.config_entries.options._progress[result["flow_id"]]
+    assert options_handler._common_handler.flow_state == {"idx": None}
+
+    # In step 1, flow state is updated with user input
+    result = await hass.config_entries.options.async_configure(
+        result["flow_id"], {"option1": "blublu"}
+    )
+    assert result["type"] == data_entry_flow.FlowResultType.FORM
+    assert result["step_id"] == "step_2"
+
+    options_handler = hass.config_entries.options._progress[result["flow_id"]]
+    assert options_handler._common_handler.flow_state == {"idx": "blublu"}
+
+    # In step 2, options were updated from flow state
+    result = await hass.config_entries.options.async_configure(
+        result["flow_id"], {"option1": "blabla"}
+    )
+    assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY
+    assert result["data"] == {
+        "idx_from_flow_state": "blublu",
+        "option1": "blabla",
+    }