Add next_flow to Repair Flows

homeassistant/components/repairs/__init__.py

@@ -7,15 +7,17 @@
 from homeassistant.helpers.typing import ConfigType
 
 from . import issue_handler, websocket_api
-from .const import DOMAIN
+from .const import DOMAIN, NextFlowType
 from .issue_handler import ConfirmRepairFlow, RepairsFlowManager
-from .models import RepairsFlow
+from .models import RepairsFlow, RepairsFlowResult
 
 __all__ = [
     "DOMAIN",
     "ConfirmRepairFlow",
+    "NextFlowType",
     "RepairsFlow",
     "RepairsFlowManager",
+    "RepairsFlowResult",
     "repairs_flow_manager",
 ]
 CONFIG_SCHEMA = cv.empty_config_schema(DOMAIN)

homeassistant/components/repairs/const.py

@@ -1,3 +1,13 @@
 """Constants for the Repairs integration."""
 
+from enum import StrEnum
+
 DOMAIN = "repairs"
+
+
+class NextFlowType(StrEnum):
+    """Extend FlowType to support additional next flow types."""
+
+    CONFIG_FLOW = "config_flow"
+    OPTIONS_FLOW = "options_flow"
+    CONFIG_SUBENTRIES_FLOW = "config_subentries_flow"

homeassistant/components/repairs/issue_handler.py

@@ -15,21 +15,21 @@
 )
 
 from .const import DOMAIN
-from .models import RepairsFlow, RepairsProtocol
+from .models import RepairsFlow, RepairsFlowResult, RepairsProtocol
 
 
 class ConfirmRepairFlow(RepairsFlow):
     """Handler for an issue fixing flow without any side effects."""
 
     async def async_step_init(
         self, user_input: dict[str, str] | None = None
-    ) -> data_entry_flow.FlowResult:
+    ) -> RepairsFlowResult:
         """Handle the first step of a fix flow."""
         return await self.async_step_confirm()
 
     async def async_step_confirm(
         self, user_input: dict[str, str] | None = None
-    ) -> data_entry_flow.FlowResult:
+    ) -> RepairsFlowResult:
         """Handle the confirm step of a fix flow."""
         if user_input is not None:
             return self.async_create_entry(data={})
@@ -46,7 +46,9 @@ async def async_step_confirm(
         )
 
 
-class RepairsFlowManager(data_entry_flow.FlowManager):
+class RepairsFlowManager(
+    data_entry_flow.FlowManager[data_entry_flow.FlowContext, RepairsFlowResult, str]
+):
     """Manage repairs flows."""
 
     async def async_create_flow(
@@ -63,7 +65,7 @@ async def async_create_flow(
         issue_registry = ir.async_get(self.hass)
         issue = issue_registry.async_get_issue(handler_key, issue_id)
         if issue is None or not issue.is_fixable:
-            raise data_entry_flow.UnknownStep
+            raise data_entry_flow.UnknownStep("Issue not found in registry")
 
         if "platforms" not in self.hass.data[DOMAIN]:
             await async_process_repairs_platforms(self.hass)
@@ -79,15 +81,19 @@ async def async_create_flow(
         flow.data = issue.data
         return flow
 
-    async def async_finish_flow(
-        self, flow: data_entry_flow.FlowHandler, result: data_entry_flow.FlowResult
-    ) -> data_entry_flow.FlowResult:
+    async def async_finish_flow(  # type: ignore[override]
+        self, flow: RepairsFlow, result: RepairsFlowResult
+    ) -> RepairsFlowResult:
         """Complete a fix flow.
 
         This method is called when a flow step returns FlowResultType.ABORT or
-        FlowResultType.CREATE_ENTRY.
+        FlowResultType.CREATE_ENTRY. When next flow is set in result, leave the issue active as the integration needs to resolve the issue
+        not the active flow.
         """
-        if result.get("type") != data_entry_flow.FlowResultType.ABORT:
+        if (
+            result.get("type") != data_entry_flow.FlowResultType.ABORT
+            and result.get("next_flow") is None
+        ):
             ir.async_delete_issue(self.hass, flow.handler, flow.init_data["issue_id"])
         return result
 

homeassistant/components/repairs/models.py

@@ -2,18 +2,103 @@
 
 from __future__ import annotations
 
-from typing import Protocol
+from collections.abc import Mapping
+from typing import Any, Protocol
 
 from homeassistant import data_entry_flow
-from homeassistant.core import HomeAssistant
+from homeassistant.config_entries import (
+    SOURCE_RECONFIGURE,
+    ConfigEntry,
+    ConfigFlowResult,
+    SubentryFlowResult,
+)
+from homeassistant.core import HomeAssistant, callback
 
+from .const import NextFlowType
 
-class RepairsFlow(data_entry_flow.FlowHandler):
+
+class RepairsFlowResult(
+    data_entry_flow.FlowResult[data_entry_flow.FlowContext, str],
+    total=False,
+):
+    """Typed context dict for repairs flow."""
+
+    next_flow: tuple[NextFlowType, str]
+    # Frontend needs this to render the dialog.
+    result: ConfigEntry
+
+
+class RepairsFlow(
+    data_entry_flow.FlowHandler[data_entry_flow.FlowContext, RepairsFlowResult, str]
+):
     """Handle a flow for fixing an issue."""
 
     issue_id: str
     data: dict[str, str | int | float | None] | None
 
+    @callback
+    def async_create_entry(
+        self,
+        *,
+        title: str | None = None,
+        data: Mapping[str, Any],
+        description: str | None = None,
+        description_placeholders: Mapping[str, str] | None = None,
+        next_flow: tuple[NextFlowType, str] | None = None,
+    ) -> RepairsFlowResult:
+        """Finish a repair flow."""
+        result: RepairsFlowResult = super().async_create_entry(
+            title=title,
+            data=data,
+            description=description,
+            description_placeholders=description_placeholders,
+        )
+
+        self._async_set_next_flow_if_valid(result, next_flow)
+
+        return result
+
+    def _async_set_next_flow_if_valid(
+        self,
+        result: RepairsFlowResult,
+        next_flow: tuple[NextFlowType, 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 not in NextFlowType:
+            raise data_entry_flow.FlowError("Invalid next_flow type")
+        flow: ConfigFlowResult | SubentryFlowResult
+        entry_id: str
+        try:
+            if flow_type in {
+                NextFlowType.CONFIG_FLOW,
+                NextFlowType.CONFIG_SUBENTRIES_FLOW,
+            }:
+                if flow_type == NextFlowType.CONFIG_FLOW:
+                    flow = self.hass.config_entries.flow.async_get(flow_id)
+                    entry_id = flow["context"]["entry_id"]
+                else:  # subentry flow
+                    flow = self.hass.config_entries.subentries.async_get(flow_id)
+                    entry_id, _ = flow["handler"]
+                if (
+                    "context" not in flow
+                    or flow["context"]["source"] != SOURCE_RECONFIGURE
+                ):  # check if reconfigure flow
+                    raise data_entry_flow.FlowError(
+                        "Next flow must be a reconfigure flow"
+                    )
+            else:
+                flow = self.hass.config_entries.options.async_get(flow_id)
+                entry_id = flow["handler"]
+        except data_entry_flow.UnknownFlow as ex:
+            raise data_entry_flow.UnknownFlow(
+                f"Unknown next flow: {flow_type}: {flow_id}"
+            ) from ex
+        result["result"] = self.hass.config_entries.async_get_known_entry(entry_id)
+        result["next_flow"] = next_flow
+
 
 class RepairsProtocol(Protocol):
     """Define the format of repairs platforms."""

homeassistant/components/repairs/websocket_api.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from http import HTTPStatus
 from typing import Any
 
@@ -13,6 +14,7 @@
 from homeassistant.components import websocket_api
 from homeassistant.components.http.data_validator import RequestDataValidator
 from homeassistant.components.http.decorators import require_admin
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import issue_registry as ir
 from homeassistant.helpers.data_entry_flow import (
@@ -21,6 +23,8 @@
 )
 
 from .const import DOMAIN
+from .issue_handler import RepairsFlowManager
+from .models import RepairsFlowResult
 
 
 @callback
@@ -107,7 +111,7 @@ def ws_list_issues(
     connection.send_result(msg["id"], {"issues": issues})
 
 
-class RepairsFlowIndexView(FlowManagerIndexView):
+class RepairsFlowIndexView(FlowManagerIndexView[RepairsFlowManager]):
     """View to create issue fix flows."""
 
     url = "/api/repairs/issues/fix"
@@ -130,19 +134,21 @@ async def post(self, request: web.Request, data: dict[str, Any]) -> web.Response
                 data["handler"],
                 data={"issue_id": data["issue_id"]},
             )
-        except data_entry_flow.UnknownHandler:
-            return self.json_message("Invalid handler specified", HTTPStatus.NOT_FOUND)
-        except data_entry_flow.UnknownStep:
-            return self.json_message(
-                "Handler does not support user", HTTPStatus.BAD_REQUEST
-            )
+        except data_entry_flow.UnknownStep as ex:
+            return self.json_message(str(ex), HTTPStatus.BAD_REQUEST)
+        except data_entry_flow.FlowError as ex:
+            return self.json_message(str(ex), HTTPStatus.BAD_REQUEST)
 
         return self.json(
             self._prepare_result_json(result),
         )
 
+    def _prepare_result_json(self, result: RepairsFlowResult) -> dict[str, Any]:  # type: ignore[override]
+        """Convert result to JSON serializable dict."""
+        return _prepare_repairs_flow_result_json(result, super()._prepare_result_json)
+
 
-class RepairsFlowResourceView(FlowManagerResourceView):
+class RepairsFlowResourceView(FlowManagerResourceView[RepairsFlowManager]):
     """View to interact with the option flow manager."""
 
     url = "/api/repairs/issues/fix/{flow_id}"
@@ -157,3 +163,23 @@ async def get(self, request: web.Request, /, flow_id: str) -> web.Response:
     async def post(self, request: web.Request, flow_id: str) -> web.Response:
         """Handle a POST request."""
         return await super().post(request, flow_id)
+
+    def _prepare_result_json(self, result: RepairsFlowResult) -> dict[str, Any]:  # type: ignore[override]
+        """Convert result to JSON serializable dict."""
+        return _prepare_repairs_flow_result_json(result, super()._prepare_result_json)
+
+
+def _prepare_repairs_flow_result_json(
+    result: RepairsFlowResult,
+    prepare_result_json: Callable[[data_entry_flow.FlowResult], dict[str, Any]],
+) -> dict[str, Any]:
+    """Convert result to JSON."""
+    if (
+        result["type"] != data_entry_flow.FlowResultType.CREATE_ENTRY
+        or "next_flow" not in result
+    ):
+        return prepare_result_json(result)
+    data = {key: val for key, val in result.items() if key not in ("data", "context")}
+    entry: ConfigEntry = result["result"]
+    data["result"] = entry.as_json_fragment
+    return data

homeassistant/config_entries.py

@@ -1760,6 +1760,7 @@ async def async_finish_flow(
             self.config_entries._async_clean_up(existing_entry)  # noqa: SLF001
 
         result["result"] = entry
+
         return result
 
     async def async_create_flow(