Fix

homeassistant/components/bayesian/config_flow.py

@@ -33,11 +33,14 @@ from homeassistant.components.update import DOMAIN as UPDATE_DOMAIN
 from homeassistant.components.weather import DOMAIN as WEATHER_DOMAIN
 from homeassistant.components.zone import DOMAIN as ZONE_DOMAIN
 from homeassistant.config_entries import (
+    SOURCE_USER,
     ConfigEntry,
     ConfigFlowResult,
     ConfigSubentry,
     ConfigSubentryData,
     ConfigSubentryFlow,
+    FlowType,
+    SubentryFlowContext,
     SubentryFlowResult,
 )
 from homeassistant.const import (
@@ -517,7 +520,23 @@ class BayesianConfigFlowHandler(SchemaConfigFlowHandler, domain=DOMAIN):
         ]
 
         self.async_config_flow_finished(data)
-        return super().async_create_entry(data=data, subentries=subentries, **kwargs)
+        return super().async_create_entry(
+            data=data,
+            subentries=subentries,
+            **kwargs,
+        )
+
+    async def async_on_create_entry(self, result: ConfigFlowResult) -> ConfigFlowResult:
+        """Create subentry flow after creating the main entry."""
+        subentry_result = await self.hass.config_entries.subentries.async_init(
+            (result["result"].entry_id, "observation"),
+            context=SubentryFlowContext(source=SOURCE_USER),
+        )
+        result["next_flow"] = (
+            FlowType.CONFIG_SUBENTRIES_FLOW,
+            subentry_result["flow_id"],
+        )
+        return result
 
 
 class ObservationSubentryFlowHandler(ConfigSubentryFlow):

homeassistant/config_entries.py

@@ -312,8 +312,8 @@ class FlowType(StrEnum):
     """Flow type."""
 
     CONFIG_FLOW = "config_flow"
-    # Add other flow types here as needed in the future,
-    # if we want to support them in the `next_flow` parameter.
+    OPTIONS_FLOW = "options_flow"
+    CONFIG_SUBENTRIES_FLOW = "config_subentries_flow"
 
 
 def _validate_item(*, disabled_by: ConfigEntryDisabler | Any | None = None) -> None:
@@ -1544,6 +1544,26 @@ class ConfigEntriesFlowManager(
                 issue_id = f"config_entry_reauth_{flow.handler}_{entry_id}"
                 ir.async_delete_issue(self.hass, HOMEASSISTANT_DOMAIN, issue_id)
 
+    def _async_validate_next_flow(
+        self,
+        result: ConfigFlowResult,
+    ) -> None:
+        """Validate and set next_flow in result if provided."""
+        if (next_flow := result.get("next_flow")) is None:
+            return
+        flow_type, flow_id = next_flow
+        if flow_type not in FlowType:
+            raise HomeAssistantError("Invalid next_flow type")
+        if flow_type == FlowType.CONFIG_FLOW:
+            # Raises UnknownFlow if the flow does not exist.
+            self.hass.config_entries.flow.async_get(flow_id)
+        if flow_type == FlowType.OPTIONS_FLOW:
+            # Raises UnknownFlow if the flow does not exist.
+            self.hass.config_entries.options.async_get(flow_id)
+        if flow_type == FlowType.CONFIG_SUBENTRIES_FLOW:
+            # Raises UnknownFlow if the flow does not exist.
+            self.hass.config_entries.subentries.async_get(flow_id)
+
     async def async_finish_flow(
         self,
         flow: data_entry_flow.FlowHandler[ConfigFlowContext, ConfigFlowResult],
@@ -1592,6 +1612,8 @@ class ConfigEntriesFlowManager(
                 self.config_entries.async_update_entry(
                     entry, discovery_keys=new_discovery_keys
                 )
+
+            self._async_validate_next_flow(result)
             return result
 
         # Mark the step as done.
@@ -1706,6 +1728,10 @@ class ConfigEntriesFlowManager(
             self.config_entries._async_clean_up(existing_entry)  # noqa: SLF001
 
         result["result"] = entry
+        if not existing_entry:
+            result = await flow.async_on_create_entry(result)
+            self._async_validate_next_flow(result)
+
         return result
 
     async def async_create_flow(
@@ -3169,21 +3195,6 @@ class ConfigFlow(ConfigEntryBaseFlow):
         """Handle a flow initialized by Zeroconf discovery."""
         return await self._async_step_discovery_without_unique_id()
 
-    def _async_set_next_flow_if_valid(
-        self,
-        result: ConfigFlowResult,
-        next_flow: tuple[FlowType, str] | None,
-    ) -> None:
-        """Validate and set next_flow in result if provided."""
-        if next_flow is None:
-            return
-        flow_type, flow_id = next_flow
-        if flow_type != FlowType.CONFIG_FLOW:
-            raise HomeAssistantError("Invalid next_flow type")
-        # Raises UnknownFlow if the flow does not exist.
-        self.hass.config_entries.flow.async_get(flow_id)
-        result["next_flow"] = next_flow
-
     @callback
     def async_abort(
         self,
@@ -3197,7 +3208,17 @@ class ConfigFlow(ConfigEntryBaseFlow):
             reason=reason,
             description_placeholders=description_placeholders,
         )
-        self._async_set_next_flow_if_valid(result, next_flow)
+        if next_flow:
+            result["next_flow"] = next_flow
+        return result
+
+    async def async_on_create_entry(self, result: ConfigFlowResult) -> ConfigFlowResult:
+        """Runs after a config flow has created a config entry.
+
+        Can be overwritten by integrations to add additional data to the result.
+        Example: creating next flow entries to the result which needs a
+        config entry created before it can start.
+        """
         return result
 
     @callback
@@ -3229,7 +3250,8 @@ class ConfigFlow(ConfigEntryBaseFlow):
         )
 
         result["minor_version"] = self.MINOR_VERSION
-        self._async_set_next_flow_if_valid(result, next_flow)
+        if next_flow:
+            result["next_flow"] = next_flow
         result["options"] = options or {}
         result["subentries"] = subentries or ()
         result["version"] = self.VERSION

tests/test_config_entries.py

@@ -2011,7 +2011,7 @@ async def test_create_entry_next_flow_invalid(
                 context={"source": config_entries.SOURCE_IMPORT},
             )
 
-        assert async_setup_entry.call_count == 0
+        assert async_setup_entry.call_count == 1
 
 
 async def test_create_entry_options(
@@ -2066,6 +2066,338 @@ async def test_create_entry_options(
         assert entries[0].options == {"example": "option"}
 
 
+async def test_on_create_entry_with_subentry_flow(
+    hass: HomeAssistant, manager: config_entries.ConfigEntries
+) -> None:
+    """Test use async_on_create_entry with creating a subentry flow."""
+
+    async def mock_async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+        """Mock setup."""
+        return True
+
+    async_setup_entry = AsyncMock(return_value=True)
+    mock_integration(
+        hass,
+        MockModule(
+            "comp",
+            async_setup=mock_async_setup,
+            async_setup_entry=async_setup_entry,
+        ),
+    )
+    mock_platform(hass, "comp.config_flow", None)
+
+    class TestFlow(config_entries.ConfigFlow):
+        """Test flow."""
+
+        @classmethod
+        @callback
+        def async_get_supported_subentry_types(
+            cls, config_entry: ConfigEntry
+        ) -> dict[str, type[config_entries.ConfigSubentryFlow]]:
+            """Return subentries supported by this integration."""
+            return {"sub_flow": TestSubentryFlowHandler}
+
+        async def async_on_create_entry(
+            self, result: config_entries.ConfigFlowResult
+        ) -> config_entries.ConfigFlowResult:
+            config_entry_id = result["result"].entry_id
+            new_flow = await hass.config_entries.subentries.async_init(
+                handler=(config_entry_id, "sub_flow"),
+                context={"source": config_entries.SOURCE_USER},
+            )
+            result["next_flow"] = (
+                config_entries.FlowType.CONFIG_SUBENTRIES_FLOW,
+                new_flow["flow_id"],
+            )
+            return result
+
+        async def async_step_user(
+            self, user_input: dict[str, Any] | None = None
+        ) -> config_entries.ConfigFlowResult:
+            """Test next step."""
+            if user_input is None:
+                return self.async_show_form(step_id="user")
+            return self.async_create_entry(
+                title="user_flow",
+                data={"flow": "user"},
+                next_flow=(config_entries.FlowType.CONFIG_FLOW, result["flow_id"]),
+            )
+
+    class TestSubentryFlowHandler(config_entries.ConfigSubentryFlow):
+        """Test subentry flow."""
+
+        async def async_step_user(
+            self, user_input: dict[str, Any] | None = None
+        ) -> config_entries.SubentryFlowResult:
+            """User flow."""
+            if user_input is None:
+                return self.async_show_form(step_id="user")
+            return self.async_create_entry(title="subentry", data={"flow": "subentry"})
+
+    with mock_config_flow("comp", TestFlow):
+        assert await async_setup_component(hass, "comp", {})
+
+        result = await hass.config_entries.flow.async_init(
+            "comp", context={"source": config_entries.SOURCE_USER}
+        )
+        await hass.async_block_till_done()
+        result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+
+        flows = hass.config_entries.flow.async_progress()
+        assert len(flows) == 0
+        sub_flows = hass.config_entries.subentries.async_progress()
+        assert len(sub_flows) == 1
+        subentry_flow = sub_flows[0]
+
+        entries = hass.config_entries.async_entries("comp")
+        assert len(entries) == 1
+        entry = entries[0]
+        assert result == {
+            "context": {"source": "user"},
+            "data": {"flow": "user"},
+            "description_placeholders": None,
+            "description": None,
+            "flow_id": ANY,
+            "handler": "comp",
+            "minor_version": 1,
+            "next_flow": (
+                config_entries.FlowType.CONFIG_SUBENTRIES_FLOW,
+                subentry_flow["flow_id"],
+            ),
+            "options": {},
+            "result": entry,
+            "subentries": (),
+            "title": "user_flow",
+            "type": FlowResultType.CREATE_ENTRY,
+            "version": 1,
+        }
+
+        result = await hass.config_entries.subentries.async_configure(
+            subentry_flow["flow_id"], {}
+        )
+        sub_flows = hass.config_entries.subentries.async_progress()
+        assert len(sub_flows) == 0
+
+        assert result == {
+            "context": {"source": "user"},
+            "data": {"flow": "subentry"},
+            "description_placeholders": None,
+            "description": None,
+            "flow_id": ANY,
+            "handler": (entry.entry_id, "sub_flow"),
+            "title": "subentry",
+            "type": FlowResultType.CREATE_ENTRY,
+            "unique_id": None,
+        }
+
+
+async def test_on_create_entry_with_options_flow(
+    hass: HomeAssistant, manager: config_entries.ConfigEntries
+) -> None:
+    """Test use async_on_create_entry with creating an options flow."""
+
+    async def mock_async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+        """Mock setup."""
+        return True
+
+    async_setup_entry = AsyncMock(return_value=True)
+    mock_integration(
+        hass,
+        MockModule(
+            "comp",
+            async_setup=mock_async_setup,
+            async_setup_entry=async_setup_entry,
+        ),
+    )
+    mock_platform(hass, "comp.config_flow", None)
+
+    class TestFlow(config_entries.ConfigFlow):
+        """Test flow."""
+
+        @staticmethod
+        @callback
+        def async_get_options_flow(config_entry: ConfigEntry) -> TestOptionsFlowHandler:
+            """Get the options flow for this handler."""
+            return TestOptionsFlowHandler()
+
+        async def async_on_create_entry(
+            self, result: config_entries.ConfigFlowResult
+        ) -> config_entries.ConfigFlowResult:
+            config_entry_id = result["result"].entry_id
+            new_flow = await hass.config_entries.options.async_init(
+                config_entry_id, context={"source": config_entries.SOURCE_USER}
+            )
+            result["next_flow"] = (
+                config_entries.FlowType.OPTIONS_FLOW,
+                new_flow["flow_id"],
+            )
+            return result
+
+        async def async_step_user(
+            self, user_input: dict[str, Any] | None = None
+        ) -> config_entries.ConfigFlowResult:
+            """Test next step."""
+            if user_input is None:
+                return self.async_show_form(step_id="user")
+            return self.async_create_entry(
+                title="user_flow",
+                data={"flow": "user"},
+                next_flow=(config_entries.FlowType.CONFIG_FLOW, result["flow_id"]),
+            )
+
+    class TestOptionsFlowHandler(config_entries.OptionsFlow):
+        """Handle Yale options."""
+
+        async def async_step_init(
+            self, user_input: dict[str, Any] | None = None
+        ) -> config_entries.ConfigFlowResult:
+            """Manage options."""
+
+            if user_input is None:
+                return self.async_show_form(step_id="init")
+            return self.async_create_entry(title="options", data={"flow": "options"})
+
+    with mock_config_flow("comp", TestFlow):
+        assert await async_setup_component(hass, "comp", {})
+
+        result = await hass.config_entries.flow.async_init(
+            "comp", context={"source": config_entries.SOURCE_USER}
+        )
+        await hass.async_block_till_done()
+        result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+
+        flows = hass.config_entries.flow.async_progress()
+        assert len(flows) == 0
+        option_flows = hass.config_entries.options.async_progress()
+        assert len(option_flows) == 1
+        option_flow = option_flows[0]
+
+        entries = hass.config_entries.async_entries("comp")
+        assert len(entries) == 1
+        entry = entries[0]
+        assert result == {
+            "context": {"source": "user"},
+            "data": {"flow": "user"},
+            "description_placeholders": None,
+            "description": None,
+            "flow_id": ANY,
+            "handler": "comp",
+            "minor_version": 1,
+            "next_flow": (
+                config_entries.FlowType.OPTIONS_FLOW,
+                option_flow["flow_id"],
+            ),
+            "options": {},
+            "result": entry,
+            "subentries": (),
+            "title": "user_flow",
+            "type": FlowResultType.CREATE_ENTRY,
+            "version": 1,
+        }
+
+        result = await hass.config_entries.options.async_configure(
+            option_flow["flow_id"], {}
+        )
+        option_flows = hass.config_entries.options.async_progress()
+        assert len(option_flows) == 0
+
+        assert result == {
+            "context": {"source": "user"},
+            "data": {"flow": "options"},
+            "description_placeholders": None,
+            "description": None,
+            "flow_id": ANY,
+            "handler": entry.entry_id,
+            "title": "options",
+            "type": FlowResultType.CREATE_ENTRY,
+        }
+
+
+@pytest.mark.parametrize(
+    ("invalid_next_flow", "error"),
+    [
+        (("invalid_flow_type", "invalid_flow_id"), HomeAssistantError),
+        ((config_entries.FlowType.CONFIG_FLOW, "invalid_flow_id"), UnknownFlow),
+        ((config_entries.FlowType.OPTIONS_FLOW, "invalid_flow_id"), UnknownFlow),
+        (
+            (config_entries.FlowType.CONFIG_SUBENTRIES_FLOW, "invalid_flow_id"),
+            UnknownFlow,
+        ),
+    ],
+)
+async def test_invalid_on_create_entry(
+    hass: HomeAssistant,
+    manager: config_entries.ConfigEntries,
+    invalid_next_flow: tuple[str, str],
+    error: type[Exception],
+) -> None:
+    """Test use invalid flows in async_on_create_entry."""
+
+    async def mock_async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+        """Mock setup."""
+        return True
+
+    async_setup_entry = AsyncMock(return_value=True)
+    mock_integration(
+        hass,
+        MockModule(
+            "comp",
+            async_setup=mock_async_setup,
+            async_setup_entry=async_setup_entry,
+        ),
+    )
+    mock_platform(hass, "comp.config_flow", None)
+
+    class TestFlow(config_entries.ConfigFlow):
+        """Test flow."""
+
+        @staticmethod
+        @callback
+        def async_get_options_flow(config_entry: ConfigEntry) -> TestOptionsFlowHandler:
+            """Get the options flow for this handler."""
+            return TestOptionsFlowHandler()
+
+        async def async_on_create_entry(
+            self, result: config_entries.ConfigFlowResult
+        ) -> config_entries.ConfigFlowResult:
+            result["next_flow"] = invalid_next_flow  # type: ignore[arg-type]
+            return result
+
+        async def async_step_user(
+            self, user_input: dict[str, Any] | None = None
+        ) -> config_entries.ConfigFlowResult:
+            """Test next step."""
+            if user_input is None:
+                return self.async_show_form(step_id="user")
+            return self.async_create_entry(
+                title="user_flow",
+                data={"flow": "user"},
+                next_flow=(config_entries.FlowType.CONFIG_FLOW, result["flow_id"]),
+            )
+
+    class TestOptionsFlowHandler(config_entries.OptionsFlow):
+        """Handle Yale options."""
+
+        async def async_step_init(
+            self, user_input: dict[str, Any] | None = None
+        ) -> config_entries.ConfigFlowResult:
+            """Manage options."""
+
+            if user_input is None:
+                return self.async_show_form(step_id="init")
+            return self.async_create_entry(title="options", data={"flow": "options"})
+
+    with mock_config_flow("comp", TestFlow):
+        assert await async_setup_component(hass, "comp", {})
+
+        result = await hass.config_entries.flow.async_init(
+            "comp", context={"source": config_entries.SOURCE_USER}
+        )
+        await hass.async_block_till_done()
+        with pytest.raises(error):
+            await hass.config_entries.flow.async_configure(result["flow_id"], {})
+
+
 async def test_entry_options(
     hass: HomeAssistant, manager: config_entries.ConfigEntries
 ) -> None: