diff --git a/homeassistant/helpers/schema_config_entry_flow.py b/homeassistant/helpers/schema_config_entry_flow.py index f992b187e9d..f29fdfe56b8 100644 --- a/homeassistant/helpers/schema_config_entry_flow.py +++ b/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) - - 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)) + return self._show_next_step(next_step_id_or_end_flow) 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( diff --git a/tests/helpers/test_schema_config_entry_flow.py b/tests/helpers/test_schema_config_entry_flow.py index de99c7b3bce..a9af60b8ec3 100644 --- a/tests/helpers/test_schema_config_entry_flow.py +++ b/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