Revert "Fix progress step recursion (#153906)" (#155866)

homeassistant/data_entry_flow.py

@@ -645,24 +645,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] | None = None
-
-    @property
-    def _progress_step_data(self) -> ProgressStepData[_FlowResultT]:
-        """Return progress step data.
-
-        A property is used instead of a simple attribute as derived classes
-        do not call super().__init__.
-        The property makes sure that the dict is initialized if needed.
-        """
-        if not self.__progress_step_data:
-            self.__progress_step_data = {
-                "tasks": {},
-                "abort_reason": "",
-                "abort_description_placeholders": MappingProxyType({}),
-                "next_step_result": None,
-            }
-        return self.__progress_step_data
+    _progress_step_data: ProgressStepData[_FlowResultT] = {
+        "tasks": {},
+        "abort_reason": "",
+        "abort_description_placeholders": MappingProxyType({}),
+        "next_step_result": None,
+    }
 
     @property
     def source(self) -> str | None:
@@ -789,10 +777,9 @@ class FlowHandler(Generic[_FlowContextT, _FlowResultT, _HandlerT]):
         self, user_input: dict[str, Any] | None = None
     ) -> _FlowResultT:
         """Abort the flow."""
-        progress_step_data = self._progress_step_data
         return self.async_abort(
-            reason=progress_step_data["abort_reason"],
-            description_placeholders=progress_step_data[
+            reason=self._progress_step_data["abort_reason"],
+            description_placeholders=self._progress_step_data[
                 "abort_description_placeholders"
             ],
         )
@@ -808,15 +795,14 @@ class FlowHandler(Generic[_FlowContextT, _FlowResultT, _HandlerT]):
         without using async_show_progress_done.
         If no next step is set, abort the flow.
         """
-        progress_step_data = self._progress_step_data
-        if (next_step_result := progress_step_data["next_step_result"]) is None:
+        if self._progress_step_data["next_step_result"] is None:
             return self.async_abort(
-                reason=progress_step_data["abort_reason"],
-                description_placeholders=progress_step_data[
+                reason=self._progress_step_data["abort_reason"],
+                description_placeholders=self._progress_step_data[
                     "abort_description_placeholders"
                 ],
             )
-        return next_step_result
+        return self._progress_step_data["next_step_result"]
 
     @callback
     def async_external_step(
@@ -1035,9 +1021,9 @@ def progress_step[
             self: FlowHandler[Any, ResultT], *args: P.args, **kwargs: P.kwargs
         ) -> ResultT:
             step_id = func.__name__.replace("async_step_", "")
-            progress_step_data = self._progress_step_data
+
             # Check if we have a progress task running
-            progress_task = progress_step_data["tasks"].get(step_id)
+            progress_task = self._progress_step_data["tasks"].get(step_id)
 
             if progress_task is None:
                 # First call - create and start the progress task
@@ -1045,30 +1031,30 @@ def progress_step[
                     func(self, *args, **kwargs),  # type: ignore[arg-type]
                     f"Progress step {step_id}",
                 )
-                progress_step_data["tasks"][step_id] = progress_task
+                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
+                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,
-                )
+                    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:
-                progress_task_result = await progress_task
+                self._progress_step_data["next_step_result"] = await progress_task
             except AbortFlow as err:
-                progress_step_data["abort_reason"] = err.reason
-                progress_step_data["abort_description_placeholders"] = (
+                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(
@@ -1076,14 +1062,7 @@ def progress_step[
                 )
             finally:
                 # Clean up task reference
-                progress_step_data["tasks"].pop(step_id, None)
-
-            # If the result type is FlowResultType.SHOW_PROGRESS_DONE
-            # an earlier show progress step has already been run and stored its result.
-            # In this case we should not overwrite the result,
-            # but just use the stored one.
-            if progress_task_result["type"] != FlowResultType.SHOW_PROGRESS_DONE:
-                progress_step_data["next_step_result"] = progress_task_result
+                self._progress_step_data["tasks"].pop(step_id, None)
 
             return self.async_show_progress_done(
                 next_step_id="_progress_step_progress_done"

tests/test_data_entry_flow.py

@@ -1,11 +1,9 @@
 """Test the flow classes."""
 
 import asyncio
-from collections.abc import Callable
 import dataclasses
 import logging
-from typing import Any
-from unittest.mock import AsyncMock, Mock, patch
+from unittest.mock import Mock, patch
 
 import pytest
 import voluptuous as vol
@@ -932,406 +930,6 @@ async def test_show_progress_fires_only_when_changed(
     )  # change (description placeholder)
 
 
-@pytest.mark.parametrize(
-    ("task_side_effect", "flow_result"),
-    [
-        (None, data_entry_flow.FlowResultType.CREATE_ENTRY),
-        (data_entry_flow.AbortFlow("fail"), data_entry_flow.FlowResultType.ABORT),
-    ],
-)
-@pytest.mark.parametrize(
-    ("description", "expected_description"),
-    [
-        (None, None),
-        ({"title": "World"}, {"title": "World"}),
-        (lambda x: {"title": "World"}, {"title": "World"}),
-    ],
-)
-async def test_progress_step(
-    hass: HomeAssistant,
-    manager: MockFlowManager,
-    description: Callable[[data_entry_flow.FlowHandler], dict[str, Any]]
-    | dict[str, Any]
-    | None,
-    expected_description: dict[str, Any] | None,
-    task_side_effect: Exception | None,
-    flow_result: data_entry_flow.FlowResultType,
-) -> None:
-    """Test progress_step decorator."""
-    manager.hass = hass
-    events = []
-    task_init_evt = asyncio.Event()
-    event_received_evt = asyncio.Event()
-    task_result = Mock()
-    task_result.side_effect = task_side_effect
-
-    @callback
-    def capture_events(event: Event) -> None:
-        events.append(event)
-        event_received_evt.set()
-
-    @manager.mock_reg_handler("test")
-    class TestFlow(data_entry_flow.FlowHandler):
-        VERSION = 5
-
-        @data_entry_flow.progress_step(description_placeholders=description)
-        async def async_step_init(self, user_input=None):
-            await task_init_evt.wait()
-            task_result()
-
-            return await self.async_step_finish()
-
-        async def async_step_finish(self, user_input=None):
-            return self.async_create_entry(data={})
-
-    hass.bus.async_listen(
-        data_entry_flow.EVENT_DATA_ENTRY_FLOW_PROGRESSED,
-        capture_events,
-    )
-
-    result = await manager.async_init("test")
-    assert result["type"] == data_entry_flow.FlowResultType.SHOW_PROGRESS
-    assert result["progress_action"] == "init"
-    description_placeholders = result["description_placeholders"]
-    assert description_placeholders == expected_description
-    assert len(manager.async_progress()) == 1
-    assert len(manager.async_progress_by_handler("test")) == 1
-    assert manager.async_get(result["flow_id"])["handler"] == "test"
-
-    # Set task one done and wait for event
-    task_init_evt.set()
-    await event_received_evt.wait()
-    event_received_evt.clear()
-    assert len(events) == 1
-    assert events[0].data == {
-        "handler": "test",
-        "flow_id": result["flow_id"],
-        "refresh": True,
-    }
-
-    # Frontend refreshes the flow
-    result = await manager.async_configure(result["flow_id"])
-    assert result["type"] == flow_result
-
-
-@pytest.mark.parametrize(
-    (
-        "task_init_side_effect",  # side effect for initial step task
-        "task_next_side_effect",  # side effect for next step task
-        "flow_result_before_init",  # result before init task is done
-        "flow_result_after_init",  # result after init task is done
-        "flow_result_after_next",  # result after next task is done
-        "flow_init_events",  # number of events fired after init task is done
-        "flow_next_events",  # number of events fired after next task is done
-        "manager_call_after_init",  # lambda to continue the flow after init task
-        "manager_call_after_next",  # lambda to continue the flow after next task
-        "before_init_task_side_effect",  # function called before init event
-        "before_next_task_side_effect",  # function called before next event
-    ),
-    [
-        (  # both steps show progress and complete successfully
-            None,
-            None,
-            data_entry_flow.FlowResultType.SHOW_PROGRESS,
-            data_entry_flow.FlowResultType.SHOW_PROGRESS,
-            data_entry_flow.FlowResultType.CREATE_ENTRY,
-            1,
-            2,
-            lambda manager, result: manager.async_configure(result["flow_id"]),
-            lambda manager, result: manager.async_configure(result["flow_id"]),
-            lambda received_event, init_task_event, next_task_event: None,
-            lambda received_event, init_task_event, next_task_event: None,
-        ),
-        (  # first step aborts
-            data_entry_flow.AbortFlow("fail"),
-            None,
-            data_entry_flow.FlowResultType.SHOW_PROGRESS,
-            data_entry_flow.FlowResultType.ABORT,
-            data_entry_flow.FlowResultType.ABORT,
-            1,
-            1,
-            lambda manager, result: manager.async_configure(result["flow_id"]),
-            lambda manager, result: AsyncMock(return_value=result)(),
-            lambda received_event, init_task_event, next_task_event: None,
-            lambda received_event, init_task_event, next_task_event: None,
-        ),
-        (  # first step shows progress, second step aborts
-            None,
-            data_entry_flow.AbortFlow("fail"),
-            data_entry_flow.FlowResultType.SHOW_PROGRESS,
-            data_entry_flow.FlowResultType.SHOW_PROGRESS,
-            data_entry_flow.FlowResultType.ABORT,
-            1,
-            2,
-            lambda manager, result: manager.async_configure(result["flow_id"]),
-            lambda manager, result: manager.async_configure(result["flow_id"]),
-            lambda received_event, init_task_event, next_task_event: None,
-            lambda received_event, init_task_event, next_task_event: None,
-        ),
-        (  # first step shows progress and second step task is already done
-            None,
-            None,
-            data_entry_flow.FlowResultType.SHOW_PROGRESS,
-            data_entry_flow.FlowResultType.CREATE_ENTRY,
-            data_entry_flow.FlowResultType.CREATE_ENTRY,
-            1,
-            1,
-            lambda manager, result: manager.async_configure(result["flow_id"]),
-            lambda manager, result: AsyncMock(return_value=result)(),
-            lambda received_event,
-            init_task_event,
-            next_task_event: next_task_event.set(),
-            lambda received_event, init_task_event, next_task_event: None,
-        ),
-        (  # both step tasks are already done and flow completes immediately
-            None,
-            None,
-            data_entry_flow.FlowResultType.SHOW_PROGRESS_DONE,
-            data_entry_flow.FlowResultType.CREATE_ENTRY,
-            data_entry_flow.FlowResultType.CREATE_ENTRY,
-            0,
-            0,
-            lambda manager, result: manager.async_configure(result["flow_id"]),
-            lambda manager, result: AsyncMock(return_value=result)(),
-            lambda received_event,
-            init_task_event,
-            next_task_event: received_event.set()
-            or init_task_event.set()
-            or next_task_event.set(),
-            lambda received_event,
-            init_task_event,
-            next_task_event: received_event.set(),
-        ),
-        (  # first step task is already done, second step shows progress and completes
-            None,
-            None,
-            data_entry_flow.FlowResultType.SHOW_PROGRESS_DONE,
-            data_entry_flow.FlowResultType.SHOW_PROGRESS,
-            data_entry_flow.FlowResultType.CREATE_ENTRY,
-            0,
-            1,
-            lambda manager, result: manager.async_configure(result["flow_id"]),
-            lambda manager, result: manager.async_configure(result["flow_id"]),
-            lambda received_event,
-            init_task_event,
-            next_task_event: received_event.set() or init_task_event.set(),
-            lambda received_event, init_task_event, next_task_event: None,
-        ),
-    ],
-)
-async def test_chaining_progress_steps(
-    hass: HomeAssistant,
-    manager: MockFlowManager,
-    task_init_side_effect: Exception | None,
-    task_next_side_effect: Exception | None,
-    flow_result_before_init: data_entry_flow.FlowResultType,
-    flow_result_after_init: data_entry_flow.FlowResultType,
-    flow_result_after_next: data_entry_flow.FlowResultType,
-    flow_init_events: int,
-    flow_next_events: int,
-    manager_call_after_init: Callable[
-        [MockFlowManager, data_entry_flow.FlowResult], Any
-    ],
-    manager_call_after_next: Callable[
-        [MockFlowManager, data_entry_flow.FlowResult], Any
-    ],
-    before_init_task_side_effect: Callable[
-        [asyncio.Event, asyncio.Event, asyncio.Event], None
-    ],
-    before_next_task_side_effect: Callable[
-        [asyncio.Event, asyncio.Event, asyncio.Event], None
-    ],
-) -> None:
-    """Test chaining two steps with progress_step decorators."""
-    manager.hass = hass
-    events = []
-    event_received_evt = asyncio.Event()
-    task_init_evt = asyncio.Event()
-    task_next_evt = asyncio.Event()
-    task_init_result = Mock()
-    task_init_result.side_effect = task_init_side_effect
-    task_next_result = Mock()
-    task_next_result.side_effect = task_next_side_effect
-
-    @callback
-    def capture_events(event: Event) -> None:
-        events.append(event)
-        event_received_evt.set()
-
-    @manager.mock_reg_handler("test")
-    class TestFlow(data_entry_flow.FlowHandler):
-        VERSION = 5
-
-        def async_remove(self) -> None:
-            # Disable event received event to allow test to finish if flow is aborted.
-            event_received_evt.set()
-
-        @data_entry_flow.progress_step()
-        async def async_step_init(self, user_input=None):
-            await task_init_evt.wait()
-            task_init_result()
-
-            return await self.async_step_next()
-
-        @data_entry_flow.progress_step()
-        async def async_step_next(self, user_input=None):
-            await task_next_evt.wait()
-            task_next_result()
-
-            return await self.async_step_finish()
-
-        async def async_step_finish(self, user_input=None):
-            return self.async_create_entry(data={})
-
-    hass.bus.async_listen(
-        data_entry_flow.EVENT_DATA_ENTRY_FLOW_PROGRESSED,
-        capture_events,
-    )
-
-    # Run side effect before first event is awaited
-    before_init_task_side_effect(event_received_evt, task_init_evt, task_next_evt)
-
-    result = await manager.async_init("test")
-    assert result["type"] == flow_result_before_init
-    assert len(manager.async_progress()) == 1
-    assert len(manager.async_progress_by_handler("test")) == 1
-    assert manager.async_get(result["flow_id"])["handler"] == "test"
-
-    # Set task init done and wait for event
-    task_init_evt.set()
-    await event_received_evt.wait()
-    event_received_evt.clear()
-    assert len(events) == flow_init_events
-
-    # Run side effect before second event is awaited
-    before_next_task_side_effect(event_received_evt, task_init_evt, task_next_evt)
-
-    # Continue the flow if needed.
-    result = await manager_call_after_init(manager, result)
-    assert result["type"] == flow_result_after_init
-
-    # Set task next done and wait for event
-    task_next_evt.set()
-    await event_received_evt.wait()
-    event_received_evt.clear()
-    assert len(events) == flow_next_events
-
-    # Continue the flow if needed.
-    result = await manager_call_after_next(manager, result)
-    assert result["type"] == flow_result_after_next
-
-
-async def test_progress_step_done_abort(
-    hass: HomeAssistant,
-    manager: MockFlowManager,
-) -> None:
-    """Test progress_step decorator without done result set."""
-    manager.hass = hass
-    events = []
-
-    @callback
-    def capture_events(event: Event) -> None:
-        events.append(event)
-
-    @manager.mock_reg_handler("test")
-    class TestFlow(data_entry_flow.FlowHandler):
-        VERSION = 5
-
-        @data_entry_flow.progress_step()
-        async def async_step_init(self, user_input=None):
-            # async_show_progress_done
-            return data_entry_flow.FlowResult(
-                flow_id=self.flow_id,
-                handler=self.handler,
-                type=data_entry_flow.FlowResultType.SHOW_PROGRESS_DONE,
-            )
-
-    hass.bus.async_listen(
-        data_entry_flow.EVENT_DATA_ENTRY_FLOW_PROGRESSED,
-        capture_events,
-    )
-
-    result = await manager.async_init("test")
-    assert result["type"] == data_entry_flow.FlowResultType.SHOW_PROGRESS_DONE
-    assert len(manager.async_progress()) == 1
-    assert len(manager.async_progress_by_handler("test")) == 1
-
-    result = await manager.async_configure(result["flow_id"])
-    assert result["type"] == data_entry_flow.FlowResultType.ABORT
-    assert result["reason"] == ""
-    assert len(manager.async_progress()) == 0
-    assert len(manager.async_progress_by_handler("test")) == 0
-    assert not events
-
-
-async def test_progress_step_result_reset(
-    hass: HomeAssistant,
-    manager: MockFlowManager,
-) -> None:
-    """Test progress_step decorator with reset result."""
-    manager.hass = hass
-    events = []
-    task_init_evt = asyncio.Event()
-    event_received_evt = asyncio.Event()
-
-    @callback
-    def capture_events(event: Event) -> None:
-        events.append(event)
-        event_received_evt.set()
-
-    @manager.mock_reg_handler("test")
-    class TestFlow(data_entry_flow.FlowHandler):
-        VERSION = 5
-
-        @data_entry_flow.progress_step()
-        async def async_step_init(self, user_input=None):
-            await task_init_evt.wait()
-
-            return await self.async_step_finish()
-
-        async def async_step_finish(self, user_input=None):
-            if user_input is None:
-                return self.async_show_form(step_id="finish")
-            return self.async_create_entry(data={})
-
-    hass.bus.async_listen(
-        data_entry_flow.EVENT_DATA_ENTRY_FLOW_PROGRESSED,
-        capture_events,
-    )
-
-    first_result = await manager.async_init("test")
-    assert first_result["type"] == data_entry_flow.FlowResultType.SHOW_PROGRESS
-    assert first_result["progress_action"] == "init"
-    assert len(manager.async_progress()) == 1
-    assert len(manager.async_progress_by_handler("test")) == 1
-    assert manager.async_get(first_result["flow_id"])["handler"] == "test"
-
-    # Set task one done and wait for event
-    task_init_evt.set()
-    await event_received_evt.wait()
-    event_received_evt.clear()
-    assert len(events) == 1
-    assert events[0].data == {
-        "handler": "test",
-        "flow_id": first_result["flow_id"],
-        "refresh": True,
-    }
-
-    # Frontend refreshes the flow
-    result = await manager.async_configure(first_result["flow_id"])
-    assert result["type"] == data_entry_flow.FlowResultType.FORM
-    assert result["step_id"] == "finish"
-
-    # Continue the flow again from the first result to test idempotency.
-    result = await manager.async_configure(first_result["flow_id"])
-    assert result["type"] == data_entry_flow.FlowResultType.FORM
-    assert result["step_id"] == "finish"
-
-    # Finish the flow
-    result = await manager.async_configure(first_result["flow_id"], {})
-    assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY
-
-
 async def test_abort_flow_exception_step(manager: MockFlowManager) -> None:
     """Test that the AbortFlow exception works in a step."""