Allow SchemaFlowFormStep.next_step to return None (#82707)

homeassistant/helpers/schema_config_entry_flow.py

@@ -53,12 +53,13 @@ class SchemaFlowFormStep(SchemaFlowStep):
     - The `validate_user_input` should raise `SchemaFlowError` is user input is invalid.
     """
 
-    next_step: Callable[[dict[str, Any]], str] | str | None = None
+    next_step: Callable[[dict[str, Any]], str | None] | str | None = None
     """Optional property to identify next step.
 
     - If `next_step` is a function, it is called if the schema validates successfully or
     if no schema is defined. The `next_step` function is passed the union of config entry
-    options and user input from previous steps.
+    options and user input from previous steps. If the function returns None, the flow is
+    ended with `FlowResultType.CREATE_ENTRY`.
     - If `next_step` is None, the flow is ended with `FlowResultType.CREATE_ENTRY`.
     """
 
@@ -147,13 +148,17 @@ class SchemaCommonFlowHandler:
     def _show_next_step_or_create_entry(
         self, form_step: SchemaFlowFormStep
     ) -> FlowResult:
-        if form_step.next_step is None:
+        next_step_id_or_end_flow: str | None
+
+        if callable(form_step.next_step):
+            next_step_id_or_end_flow = form_step.next_step(self._options)
+        else:
+            next_step_id_or_end_flow = form_step.next_step
+
+        if next_step_id_or_end_flow is None:
             # Flow done, create entry or update config entry options
             return self._handler.async_create_entry(data=self._options)
-
+        return self._show_next_step(next_step_id_or_end_flow)
-        if isinstance(form_step.next_step, str):
-            return self._show_next_step(form_step.next_step)
-        return self._show_next_step(form_step.next_step(self._options))
 
     def _show_next_step(
         self,
@@ -203,11 +208,14 @@ class SchemaCommonFlowHandler:
         errors = {"base": str(error)} if error else None
 
         # Show form for next step
+        last_step = None
+        if not callable(form_step.next_step):
+            last_step = form_step.next_step is None
         return self._handler.async_show_form(
             step_id=next_step_id,
             data_schema=data_schema,
             errors=errors,
-            last_step=form_step.next_step is None,
+            last_step=last_step,
         )
 
     async def _async_menu_step(

tests/helpers/test_schema_config_entry_flow.py

@@ -375,3 +375,71 @@ async def test_schema_none(hass: HomeAssistant) -> None:
 
         result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
         assert result["type"] == FlowResultType.CREATE_ENTRY
+
+
+async def test_last_step(hass: HomeAssistant) -> None:
+    """Test SchemaFlowFormStep with schema set to None."""
+
+    CONFIG_FLOW: dict[str, SchemaFlowFormStep | SchemaFlowMenuStep] = {
+        "user": SchemaFlowFormStep(next_step="step1"),
+        "step1": SchemaFlowFormStep(vol.Schema({}), next_step="step2"),
+        "step2": SchemaFlowFormStep(vol.Schema({}), next_step=lambda _: "step3"),
+        "step3": SchemaFlowFormStep(vol.Schema({}), next_step=None),
+    }
+
+    class TestConfigFlow(SchemaConfigFlowHandler, domain=TEST_DOMAIN):
+        """Handle a config or options flow for Derivative."""
+
+        config_flow = CONFIG_FLOW
+
+    mock_platform(hass, f"{TEST_DOMAIN}.config_flow")
+    with patch.dict(config_entries.HANDLERS, {TEST_DOMAIN: TestConfigFlow}):
+        result = await hass.config_entries.flow.async_init(
+            TEST_DOMAIN, context={"source": "user"}
+        )
+        assert result["type"] == FlowResultType.FORM
+        assert result["step_id"] == "step1"
+        assert result["last_step"] is False
+
+        result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+        assert result["type"] == FlowResultType.FORM
+        assert result["step_id"] == "step2"
+        assert result["last_step"] is None
+
+        result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+        assert result["type"] == FlowResultType.FORM
+        assert result["step_id"] == "step3"
+        assert result["last_step"] is True
+
+        result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+        assert result["type"] == FlowResultType.CREATE_ENTRY
+
+
+async def test_next_step_function(hass: HomeAssistant) -> None:
+    """Test SchemaFlowFormStep with a next_step function."""
+
+    CONFIG_FLOW: dict[str, SchemaFlowFormStep | SchemaFlowMenuStep] = {
+        "user": SchemaFlowFormStep(next_step="step1"),
+        "step1": SchemaFlowFormStep(vol.Schema({}), next_step=lambda _: "step2"),
+        "step2": SchemaFlowFormStep(vol.Schema({}), next_step=lambda _: None),
+    }
+
+    class TestConfigFlow(SchemaConfigFlowHandler, domain=TEST_DOMAIN):
+        """Handle a config or options flow for Derivative."""
+
+        config_flow = CONFIG_FLOW
+
+    mock_platform(hass, f"{TEST_DOMAIN}.config_flow")
+    with patch.dict(config_entries.HANDLERS, {TEST_DOMAIN: TestConfigFlow}):
+        result = await hass.config_entries.flow.async_init(
+            TEST_DOMAIN, context={"source": "user"}
+        )
+        assert result["type"] == FlowResultType.FORM
+        assert result["step_id"] == "step1"
+
+        result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+        assert result["type"] == FlowResultType.FORM
+        assert result["step_id"] == "step2"
+
+        result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+        assert result["type"] == FlowResultType.CREATE_ENTRY