From fbb6c1d0f0a025cb42fe086428169e108fd2d195 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Thu, 11 Jan 2024 13:08:54 +0100
Subject: [PATCH] Improve calls to async_show_progress in matter (#107791)

---
 .../components/matter/config_flow.py          | 71 +++++++++----------
 1 file changed, 33 insertions(+), 38 deletions(-)

diff --git a/homeassistant/components/matter/config_flow.py b/homeassistant/components/matter/config_flow.py
index 6e25370d86a..1636790c4cb 100644
--- a/homeassistant/components/matter/config_flow.py
+++ b/homeassistant/components/matter/config_flow.py
@@ -76,19 +76,23 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         """Install Matter Server add-on."""
         if not self.install_task:
             self.install_task = self.hass.async_create_task(self._async_install_addon())
+
+        if not self.install_task.done():
             return self.async_show_progress(
-                step_id="install_addon", progress_action="install_addon"
+                step_id="install_addon",
+                progress_action="install_addon",
+                progress_task=self.install_task,
             )
 
         try:
             await self.install_task
         except AddonError as err:
-            self.install_task = None
             LOGGER.error(err)
             return self.async_show_progress_done(next_step_id="install_failed")
+        finally:
+            self.install_task = None
 
         self.integration_created_addon = True
-        self.install_task = None
 
         return self.async_show_progress_done(next_step_id="start_addon")
 
@@ -101,13 +105,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     async def _async_install_addon(self) -> None:
         """Install the Matter Server add-on."""
         addon_manager: AddonManager = get_addon_manager(self.hass)
-        try:
-            await addon_manager.async_schedule_install_addon()
-        finally:
-            # Continue the flow after show progress when the task is done.
-            self.hass.async_create_task(
-                self.hass.config_entries.flow.async_configure(flow_id=self.flow_id)
-            )
+        await addon_manager.async_schedule_install_addon()
 
     async def _async_get_addon_discovery_info(self) -> dict:
         """Return add-on discovery info."""
@@ -126,18 +124,21 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         """Start Matter Server add-on."""
         if not self.start_task:
             self.start_task = self.hass.async_create_task(self._async_start_addon())
+        if not self.start_task.done():
             return self.async_show_progress(
-                step_id="start_addon", progress_action="start_addon"
+                step_id="start_addon",
+                progress_action="start_addon",
+                progress_task=self.start_task,
             )
 
         try:
             await self.start_task
         except (FailedConnect, AddonError, AbortFlow) as err:
-            self.start_task = None
             LOGGER.error(err)
             return self.async_show_progress_done(next_step_id="start_failed")
+        finally:
+            self.start_task = None
 
-        self.start_task = None
         return self.async_show_progress_done(next_step_id="finish_addon_setup")
 
     async def async_step_start_failed(
@@ -150,33 +151,27 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         """Start the Matter Server add-on."""
         addon_manager: AddonManager = get_addon_manager(self.hass)
 
-        try:
-            await addon_manager.async_schedule_start_addon()
-            # Sleep some seconds to let the add-on start properly before connecting.
-            for _ in range(ADDON_SETUP_TIMEOUT_ROUNDS):
-                await asyncio.sleep(ADDON_SETUP_TIMEOUT)
-                try:
-                    if not (ws_address := self.ws_address):
-                        discovery_info = await self._async_get_addon_discovery_info()
-                        ws_address = self.ws_address = build_ws_address(
-                            discovery_info["host"], discovery_info["port"]
-                        )
-                    await validate_input(self.hass, {CONF_URL: ws_address})
-                except (AbortFlow, CannotConnect) as err:
-                    LOGGER.debug(
-                        "Add-on not ready yet, waiting %s seconds: %s",
-                        ADDON_SETUP_TIMEOUT,
-                        err,
+        await addon_manager.async_schedule_start_addon()
+        # Sleep some seconds to let the add-on start properly before connecting.
+        for _ in range(ADDON_SETUP_TIMEOUT_ROUNDS):
+            await asyncio.sleep(ADDON_SETUP_TIMEOUT)
+            try:
+                if not (ws_address := self.ws_address):
+                    discovery_info = await self._async_get_addon_discovery_info()
+                    ws_address = self.ws_address = build_ws_address(
+                        discovery_info["host"], discovery_info["port"]
                     )
-                else:
-                    break
+                await validate_input(self.hass, {CONF_URL: ws_address})
+            except (AbortFlow, CannotConnect) as err:
+                LOGGER.debug(
+                    "Add-on not ready yet, waiting %s seconds: %s",
+                    ADDON_SETUP_TIMEOUT,
+                    err,
+                )
             else:
-                raise FailedConnect("Failed to start Matter Server add-on: timeout")
-        finally:
-            # Continue the flow after show progress when the task is done.
-            self.hass.async_create_task(
-                self.hass.config_entries.flow.async_configure(flow_id=self.flow_id)
-            )
+                break
+        else:
+            raise FailedConnect("Failed to start Matter Server add-on: timeout")
 
     async def _async_get_addon_info(self) -> AddonInfo:
         """Return Matter Server add-on info."""