Move first probe firmware to firmware progress in hardware flow (#152819)

homeassistant/components/homeassistant_hardware/firmware_config_flow.py

@@ -88,7 +88,7 @@ class BaseFirmwareInstallFlow(ConfigEntryBaseFlow, ABC):
         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
+        self.firmware_install_task: asyncio.Task[None] | None = None
         self.installing_firmware_name: str | None = None
 
     def _get_translation_placeholders(self) -> dict[str, str]:
@@ -184,9 +184,58 @@ class BaseFirmwareInstallFlow(ConfigEntryBaseFlow, ABC):
         step_id: str,
         next_step_id: str,
     ) -> ConfigFlowResult:
+        """Show progress dialog for installing firmware."""
+        if not self.firmware_install_task:
+            self.firmware_install_task = self.hass.async_create_task(
+                self._install_firmware(
+                    fw_update_url,
+                    fw_type,
+                    firmware_name,
+                    expected_installed_firmware_type,
+                ),
+                f"Install {firmware_name} firmware",
+            )
+        if not self.firmware_install_task.done():
+            return self.async_show_progress(
+                step_id=step_id,
+                progress_action="install_firmware",
+                description_placeholders={
+                    **self._get_translation_placeholders(),
+                    "firmware_name": firmware_name,
+                },
+                progress_task=self.firmware_install_task,
+            )
+
+        try:
+            await self.firmware_install_task
+        except AbortFlow as err:
+            return self.async_show_progress_done(
+                next_step_id=err.reason,
+            )
+        except HomeAssistantError:
+            _LOGGER.exception("Failed to flash firmware")
+            return self.async_show_progress_done(next_step_id="firmware_install_failed")
+        finally:
+            self.firmware_install_task = None
+
+        return self.async_show_progress_done(next_step_id=next_step_id)
+
+    async def _install_firmware(
+        self,
+        fw_update_url: str,
+        fw_type: str,
+        firmware_name: str,
+        expected_installed_firmware_type: ApplicationType,
+    ) -> None:
+        """Install firmware."""
+        if not await self._probe_firmware_info():
+            raise AbortFlow(
+                reason="unsupported_firmware",
+                description_placeholders=self._get_translation_placeholders(),
+            )
+
         assert self._device is not None
 
-        if not self.firmware_install_task:
         # Keep track of the firmware we're working with, for error messages
         self.installing_firmware_name = firmware_name
 
@@ -194,8 +243,7 @@ class BaseFirmwareInstallFlow(ConfigEntryBaseFlow, ABC):
         # installed: upgrading to the latest release of the current firmware type
         # isn't strictly necessary for functionality.
         firmware_install_required = self._probed_firmware_info is None or (
-                self._probed_firmware_info.firmware_type
+            self._probed_firmware_info.firmware_type != expected_installed_firmware_type
-                != expected_installed_firmware_type
         )
 
         session = async_get_clientsession(self.hass)
@@ -206,21 +254,15 @@ class BaseFirmwareInstallFlow(ConfigEntryBaseFlow, ABC):
             fw_manifest = next(
                 fw for fw in manifest.firmwares if fw.filename.startswith(fw_type)
             )
-            except (StopIteration, TimeoutError, ClientError, ManifestMissing):
+        except (StopIteration, TimeoutError, ClientError, ManifestMissing) as err:
-                _LOGGER.warning(
+            _LOGGER.warning("Failed to fetch firmware update manifest", exc_info=True)
-                    "Failed to fetch firmware update manifest", exc_info=True
-                )
 
             # Not having internet access should not prevent setup
             if not firmware_install_required:
-                    _LOGGER.debug(
+                _LOGGER.debug("Skipping firmware upgrade due to index download failure")
-                        "Skipping firmware upgrade due to index download failure"
+                return
-                    )
-                    return self.async_show_progress_done(next_step_id=next_step_id)
 
-                return self.async_show_progress_done(
+            raise AbortFlow(reason="firmware_download_failed") from err
-                    next_step_id="firmware_download_failed"
-                )
 
         if not firmware_install_required:
             assert self._probed_firmware_info is not None
@@ -236,27 +278,22 @@ class BaseFirmwareInstallFlow(ConfigEntryBaseFlow, ABC):
                     probed_fw_version,
                     fw_version,
                 )
-                    return self.async_show_progress_done(next_step_id=next_step_id)
+                return
 
         try:
             fw_data = await client.async_fetch_firmware(fw_manifest)
-            except (TimeoutError, ClientError, ValueError):
+        except (TimeoutError, ClientError, ValueError) as err:
             _LOGGER.warning("Failed to fetch firmware update", exc_info=True)
 
             # If we cannot download new firmware, we shouldn't block setup
             if not firmware_install_required:
-                    _LOGGER.debug(
+                _LOGGER.debug("Skipping firmware upgrade due to image download failure")
-                        "Skipping firmware upgrade due to image download failure"
+                return
-                    )
-                    return self.async_show_progress_done(next_step_id=next_step_id)
 
             # Otherwise, fail
-                return self.async_show_progress_done(
+            raise AbortFlow(reason="firmware_download_failed") from err
-                    next_step_id="firmware_download_failed"
-                )
 
-            self.firmware_install_task = self.hass.async_create_task(
+        await async_flash_silabs_firmware(
-                async_flash_silabs_firmware(
             hass=self.hass,
             device=self._device,
             fw_data=fw_data,
@@ -265,29 +302,8 @@ class BaseFirmwareInstallFlow(ConfigEntryBaseFlow, ABC):
             progress_callback=lambda offset, total: self.async_update_progress(
                 offset / total
             ),
-                ),
-                f"Flash {firmware_name} firmware",
         )
 
-        if not self.firmware_install_task.done():
-            return self.async_show_progress(
-                step_id=step_id,
-                progress_action="install_firmware",
-                description_placeholders={
-                    **self._get_translation_placeholders(),
-                    "firmware_name": firmware_name,
-                },
-                progress_task=self.firmware_install_task,
-            )
-
-        try:
-            await self.firmware_install_task
-        except HomeAssistantError:
-            _LOGGER.exception("Failed to flash firmware")
-            return self.async_show_progress_done(next_step_id="firmware_install_failed")
-
-        return self.async_show_progress_done(next_step_id=next_step_id)
-
     async def _configure_and_start_otbr_addon(self) -> None:
         """Configure and start the OTBR addon."""
 
@@ -353,6 +369,15 @@ class BaseFirmwareInstallFlow(ConfigEntryBaseFlow, ABC):
             },
         )
 
+    async def async_step_unsupported_firmware(
+        self, user_input: dict[str, Any] | None = None
+    ) -> ConfigFlowResult:
+        """Abort when unsupported firmware is detected."""
+        return self.async_abort(
+            reason="unsupported_firmware",
+            description_placeholders=self._get_translation_placeholders(),
+        )
+
     async def async_step_zigbee_installation_type(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
@@ -406,12 +431,6 @@ class BaseFirmwareInstallFlow(ConfigEntryBaseFlow, ABC):
 
     async def _async_continue_picked_firmware(self) -> ConfigFlowResult:
         """Continue to the picked firmware step."""
-        if not await self._probe_firmware_info():
-            return self.async_abort(
-                reason="unsupported_firmware",
-                description_placeholders=self._get_translation_placeholders(),
-            )
-
         if self._picked_firmware_type == PickedFirmwareType.ZIGBEE:
             return await self.async_step_install_zigbee_firmware()