Add progress step decorator for easier config flows (home-assistant#152739)

MindFreeze · cdce8p · MartinHjelmare · web-flow · commit dfbaf66021c5 · 2025-09-24T18:18:42.000+03:00
Co-authored-by: Marc Mueller &lt;30130371+cdce8p@users.noreply.github.com&gt;
Co-authored-by: Martin Hjelmare &lt;marhje52@gmail.com&gt;
diff --git a/homeassistant/components/homeassistant_hardware/firmware_config_flow.py b/homeassistant/components/homeassistant_hardware/firmware_config_flow.py
@@ -28,7 +28,7 @@
     OptionsFlow,
 )
 from homeassistant.core import callback
-from homeassistant.data_entry_flow import AbortFlow
+from homeassistant.data_entry_flow import AbortFlow, progress_step
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.hassio import is_hassio
@@ -72,8 +72,6 @@ class BaseFirmwareInstallFlow(ConfigEntryBaseFlow, ABC):
     """Base flow to install firmware."""
 
     ZIGBEE_BAUDRATE = 115200  # Default, subclasses may override
-    _failed_addon_name: str
-    _failed_addon_reason: str
     _picked_firmware_type: PickedFirmwareType
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
@@ -85,8 +83,6 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self._hardware_name: str = "unknown"  # To be set in a subclass
         self._zigbee_integration = ZigbeeIntegration.ZHA
 
-        self.addon_install_task: asyncio.Task | None = None
-        self.addon_start_task: asyncio.Task | None = None
         self.addon_uninstall_task: asyncio.Task | None = None
         self.firmware_install_task: asyncio.Task[None] | None = None
         self.installing_firmware_name: str | None = None
@@ -486,18 +482,6 @@ async def async_step_install_zigbee_firmware(
         """Install Zigbee firmware."""
         raise NotImplementedError
 
-    async def async_step_addon_operation_failed(
-        self, user_input: dict[str, Any] | None = None
-    ) -> ConfigFlowResult:
-        """Abort when add-on installation or start failed."""
-        return self.async_abort(
-            reason=self._failed_addon_reason,
-            description_placeholders={
-                **self._get_translation_placeholders(),
-                "addon_name": self._failed_addon_name,
-            },
-        )
-
     async def async_step_pre_confirm_zigbee(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
@@ -561,6 +545,12 @@ async def async_step_install_thread_firmware(
         """Install Thread firmware."""
         raise NotImplementedError
 
+    @progress_step(
+        description_placeholders=lambda self: {
+            **self._get_translation_placeholders(),
+            "addon_name": get_otbr_addon_manager(self.hass).addon_name,
+        }
+    )
     async def async_step_install_otbr_addon(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
@@ -570,70 +560,43 @@ async def async_step_install_otbr_addon(
 
         _LOGGER.debug("OTBR addon info: %s", addon_info)
 
-        if not self.addon_install_task:
-            self.addon_install_task = self.hass.async_create_task(
-                addon_manager.async_install_addon_waiting(),
-                "OTBR addon install",
-            )
-
-        if not self.addon_install_task.done():
-            return self.async_show_progress(
-                step_id="install_otbr_addon",
-                progress_action="install_addon",
+        try:
+            await addon_manager.async_install_addon_waiting()
+        except AddonError as err:
+            _LOGGER.error(err)
+            raise AbortFlow(
+                "addon_install_failed",
                 description_placeholders={
                     **self._get_translation_placeholders(),
                     "addon_name": addon_manager.addon_name,
                 },
-                progress_task=self.addon_install_task,
-            )
-
-        try:
-            await self.addon_install_task
-        except AddonError as err:
-            _LOGGER.error(err)
-            self._failed_addon_name = addon_manager.addon_name
-            self._failed_addon_reason = "addon_install_failed"
-            return self.async_show_progress_done(next_step_id="addon_operation_failed")
-        finally:
-            self.addon_install_task = None
+            ) from err
 
-        return self.async_show_progress_done(next_step_id="finish_thread_installation")
+        return await self.async_step_finish_thread_installation()
 
+    @progress_step(
+        description_placeholders=lambda self: {
+            **self._get_translation_placeholders(),
+            "addon_name": get_otbr_addon_manager(self.hass).addon_name,
+        }
+    )
     async def async_step_start_otbr_addon(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Configure OTBR to point to the SkyConnect and run the addon."""
-        otbr_manager = get_otbr_addon_manager(self.hass)
-
-        if not self.addon_start_task:
-            self.addon_start_task = self.hass.async_create_task(
-                self._configure_and_start_otbr_addon()
-            )
-
-        if not self.addon_start_task.done():
-            return self.async_show_progress(
-                step_id="start_otbr_addon",
-                progress_action="start_otbr_addon",
+        try:
+            await self._configure_and_start_otbr_addon()
+        except AddonError as err:
+            _LOGGER.error(err)
+            raise AbortFlow(
+                "addon_start_failed",
                 description_placeholders={
                     **self._get_translation_placeholders(),
-                    "addon_name": otbr_manager.addon_name,
+                    "addon_name": get_otbr_addon_manager(self.hass).addon_name,
                 },
-                progress_task=self.addon_start_task,
-            )
-
-        try:
-            await self.addon_start_task
-        except (AddonError, AbortFlow) as err:
-            _LOGGER.error(err)
-            self._failed_addon_name = otbr_manager.addon_name
-            self._failed_addon_reason = (
-                err.reason if isinstance(err, AbortFlow) else "addon_start_failed"
-            )
-            return self.async_show_progress_done(next_step_id="addon_operation_failed")
-        finally:
-            self.addon_start_task = None
+            ) from err
 
-        return self.async_show_progress_done(next_step_id="pre_confirm_otbr")
+        return await self.async_step_pre_confirm_otbr()
 
     async def async_step_pre_confirm_otbr(
         self, user_input: dict[str, Any] | None = None
diff --git a/homeassistant/data_entry_flow.py b/homeassistant/data_entry_flow.py
@@ -5,14 +5,15 @@
 import abc
 import asyncio
 from collections import defaultdict
-from collections.abc import Callable, Container, Hashable, Iterable, Mapping
+from collections.abc import Callable, Container, Coroutine, Hashable, Iterable, Mapping
 from contextlib import suppress
 import copy
 from dataclasses import dataclass
 from enum import StrEnum
+import functools
 import logging
 from types import MappingProxyType
-from typing import Any, Generic, Required, TypedDict, TypeVar, cast
+from typing import Any, Concatenate, Generic, Required, TypedDict, TypeVar, cast
 
 import voluptuous as vol
 
@@ -150,6 +151,15 @@ class FlowResult(TypedDict, Generic[_FlowContextT, _HandlerT], total=False):
     url: str
 
 
+class ProgressStepData[_FlowResultT](TypedDict):
+    """Typed data for progress step tracking."""
+
+    tasks: dict[str, asyncio.Task[Any]]
+    abort_reason: str
+    abort_description_placeholders: Mapping[str, str]
+    next_step_result: _FlowResultT | None
+
+
 def _map_error_to_schema_errors(
     schema_errors: dict[str, Any],
     error: vol.Invalid,
@@ -639,6 +649,12 @@ class FlowHandler(Generic[_FlowContextT, _FlowResultT, _HandlerT]):
     __progress_task: asyncio.Task[Any] | None = None
     __no_progress_task_reported = False
     deprecated_show_progress = False
+    _progress_step_data: ProgressStepData[_FlowResultT] = {
+        "tasks": {},
+        "abort_reason": "",
+        "abort_description_placeholders": MappingProxyType({}),
+        "next_step_result": None,
+    }
 
     @property
     def source(self) -> str | None:
@@ -761,6 +777,37 @@ def async_abort(
             description_placeholders=description_placeholders,
         )
 
+    async def async_step__progress_step_abort(
+        self, user_input: dict[str, Any] | None = None
+    ) -> _FlowResultT:
+        """Abort the flow."""
+        return self.async_abort(
+            reason=self._progress_step_data["abort_reason"],
+            description_placeholders=self._progress_step_data[
+                "abort_description_placeholders"
+            ],
+        )
+
+    async def async_step__progress_step_progress_done(
+        self, user_input: dict[str, Any] | None = None
+    ) -> _FlowResultT:
+        """Progress done. Return the next step.
+
+        Used by the progress_step decorator
+        to allow decorated step methods
+        to call the next step method, to change step,
+        without using async_show_progress_done.
+        If no next step is set, abort the flow.
+        """
+        if self._progress_step_data["next_step_result"] is None:
+            return self.async_abort(
+                reason=self._progress_step_data["abort_reason"],
+                description_placeholders=self._progress_step_data[
+                    "abort_description_placeholders"
+                ],
+            )
+        return self._progress_step_data["next_step_result"]
+
     @callback
     def async_external_step(
         self,
@@ -930,3 +977,90 @@ def __init__(
     def __call__(self, value: Any) -> Any:
         """Validate input."""
         return self.schema(value)
+
+
+type _FuncType[_T: FlowHandler[Any, Any, Any], _R: FlowResult[Any, Any], **_P] = (
+    Callable[Concatenate[_T, _P], Coroutine[Any, Any, _R]]
+)
+
+
+def progress_step[
+    HandlerT: FlowHandler[Any, Any, Any],
+    ResultT: FlowResult[Any, Any],
+    **P,
+](
+    description_placeholders: (
+        dict[str, str] | Callable[[Any], dict[str, str]] | None
+    ) = None,
+) -> Callable[[_FuncType[HandlerT, ResultT, P]], _FuncType[HandlerT, ResultT, P]]:
+    """Decorator to create a progress step from an async function.
+
+    The decorated method should be a step method
+    which needs to show progress.
+    The method should accept dict[str, Any] as user_input
+    and should return a FlowResult or raise AbortFlow.
+    The method can call self.async_update_progress(progress)
+    to update progress.
+
+    Args:
+        description_placeholders: Static dict or callable that returns dict for progress UI placeholders.
+    """
+
+    def decorator(
+        func: _FuncType[HandlerT, ResultT, P],
+    ) -> _FuncType[HandlerT, ResultT, P]:
+        @functools.wraps(func)
+        async def wrapper(
+            self: FlowHandler[Any, ResultT], *args: P.args, **kwargs: P.kwargs
+        ) -> ResultT:
+            step_id = func.__name__.replace("async_step_", "")
+
+            # Check if we have a progress task running
+            progress_task = self._progress_step_data["tasks"].get(step_id)
+
+            if progress_task is None:
+                # First call - create and start the progress task
+                progress_task = self.hass.async_create_task(
+                    func(self, *args, **kwargs),  # type: ignore[arg-type]
+                    f"Progress step {step_id}",
+                )
+                self._progress_step_data["tasks"][step_id] = progress_task
+
+                if not progress_task.done():
+                    # Handle description placeholders
+                    placeholders = None
+                    if description_placeholders is not None:
+                        if callable(description_placeholders):
+                            placeholders = description_placeholders(self)
+                        else:
+                            placeholders = description_placeholders
+
+                    return self.async_show_progress(
+                        step_id=step_id,
+                        progress_action=step_id,
+                        progress_task=progress_task,
+                        description_placeholders=placeholders,
+                    )
+
+            # Task is done or this is a subsequent call
+            try:
+                self._progress_step_data["next_step_result"] = await progress_task
+            except AbortFlow as err:
+                self._progress_step_data["abort_reason"] = err.reason
+                self._progress_step_data["abort_description_placeholders"] = (
+                    err.description_placeholders or {}
+                )
+                return self.async_show_progress_done(
+                    next_step_id="_progress_step_abort"
+                )
+            finally:
+                # Clean up task reference
+                self._progress_step_data["tasks"].pop(step_id, None)
+
+            return self.async_show_progress_done(
+                next_step_id="_progress_step_progress_done"
+            )
+
+        return wrapper
+
+    return decorator
diff --git a/tests/components/homeassistant_hardware/test_config_flow.py b/tests/components/homeassistant_hardware/test_config_flow.py
@@ -844,7 +844,7 @@ async def test_options_flow_zigbee_to_thread(
 
         assert result["type"] is FlowResultType.SHOW_PROGRESS
         assert result["step_id"] == "install_otbr_addon"
-        assert result["progress_action"] == "install_addon"
+        assert result["progress_action"] == "install_otbr_addon"
 
         await hass.async_block_till_done(wait_background_tasks=True)
 
