Refactor zwave_js config flow (#51720)

homeassistant/components/zwave_js/config_flow.py

@@ -1,6 +1,7 @@
 """Config flow for Z-Wave JS integration."""
 from __future__ import annotations
 
+from abc import abstractmethod
 import asyncio
 import logging
 from typing import Any
@@ -14,7 +15,12 @@ from homeassistant import config_entries, exceptions
 from homeassistant.components.hassio import is_hassio
 from homeassistant.const import CONF_URL
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.data_entry_flow import AbortFlow, FlowResult
+from homeassistant.data_entry_flow import (
+    AbortFlow,
+    FlowHandler,
+    FlowManager,
+    FlowResult,
+)
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
 from .addon import AddonError, AddonInfo, AddonManager, AddonState, get_addon_manager
@@ -38,7 +44,12 @@ ADDON_SETUP_TIMEOUT_ROUNDS = 4
 SERVER_VERSION_TIMEOUT = 10
 
 ON_SUPERVISOR_SCHEMA = vol.Schema({vol.Optional(CONF_USE_ADDON, default=True): bool})
-STEP_USER_DATA_SCHEMA = vol.Schema({vol.Required(CONF_URL, default=DEFAULT_URL): str})
+
+
+def get_manual_schema(user_input: dict[str, Any]) -> vol.Schema:
+    """Return a schema for the manual step."""
+    default_url = user_input.get(CONF_URL, DEFAULT_URL)
+    return vol.Schema({vol.Required(CONF_URL, default=default_url): str})
 
 
 async def validate_input(hass: HomeAssistant, user_input: dict) -> VersionInfo:
@@ -70,135 +81,24 @@ async def async_get_version_info(hass: HomeAssistant, ws_address: str) -> Versio
     return version_info
 
 
-class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
+class BaseZwaveJSFlow(FlowHandler):
-    """Handle a config flow for Z-Wave JS."""
+    """Represent the base config flow for Z-Wave JS."""
-
-    VERSION = 1
 
     def __init__(self) -> None:
         """Set up flow instance."""
         self.network_key: str | None = None
         self.usb_path: str | None = None
-        self.use_addon = False
         self.ws_address: str | None = None
         # If we install the add-on we should uninstall it on entry remove.
         self.integration_created_addon = False
         self.install_task: asyncio.Task | None = None
         self.start_task: asyncio.Task | None = None
+        self.version_info: VersionInfo | None = None
 
-    async def async_step_user(
+    @property
-        self, user_input: dict[str, Any] | None = None
+    @abstractmethod
-    ) -> FlowResult:
+    def flow_manager(self) -> FlowManager:
-        """Handle the initial step."""
+        """Return the flow manager of the flow."""
-        if is_hassio(self.hass):
-            return await self.async_step_on_supervisor()
-
-        return await self.async_step_manual()
-
-    async def async_step_manual(
-        self, user_input: dict[str, Any] | None = None
-    ) -> FlowResult:
-        """Handle a manual configuration."""
-        if user_input is None:
-            return self.async_show_form(
-                step_id="manual", data_schema=STEP_USER_DATA_SCHEMA
-            )
-
-        errors = {}
-
-        try:
-            version_info = await validate_input(self.hass, user_input)
-        except InvalidInput as err:
-            errors["base"] = err.error
-        except Exception:  # pylint: disable=broad-except
-            _LOGGER.exception("Unexpected exception")
-            errors["base"] = "unknown"
-        else:
-            await self.async_set_unique_id(
-                version_info.home_id, raise_on_progress=False
-            )
-            # Make sure we disable any add-on handling
-            # if the controller is reconfigured in a manual step.
-            self._abort_if_unique_id_configured(
-                updates={
-                    **user_input,
-                    CONF_USE_ADDON: False,
-                    CONF_INTEGRATION_CREATED_ADDON: False,
-                }
-            )
-            self.ws_address = user_input[CONF_URL]
-            return self._async_create_entry_from_vars()
-
-        return self.async_show_form(
-            step_id="manual", data_schema=STEP_USER_DATA_SCHEMA, errors=errors
-        )
-
-    async def async_step_hassio(self, discovery_info: dict[str, Any]) -> FlowResult:
-        """Receive configuration from add-on discovery info.
-
-        This flow is triggered by the Z-Wave JS add-on.
-        """
-        self.ws_address = f"ws://{discovery_info['host']}:{discovery_info['port']}"
-        try:
-            version_info = await async_get_version_info(self.hass, self.ws_address)
-        except CannotConnect:
-            return self.async_abort(reason="cannot_connect")
-
-        await self.async_set_unique_id(version_info.home_id)
-        self._abort_if_unique_id_configured(updates={CONF_URL: self.ws_address})
-
-        return await self.async_step_hassio_confirm()
-
-    async def async_step_hassio_confirm(
-        self, user_input: dict[str, Any] | None = None
-    ) -> FlowResult:
-        """Confirm the add-on discovery."""
-        if user_input is not None:
-            return await self.async_step_on_supervisor(
-                user_input={CONF_USE_ADDON: True}
-            )
-
-        return self.async_show_form(step_id="hassio_confirm")
-
-    @callback
-    def _async_create_entry_from_vars(self) -> FlowResult:
-        """Return a config entry for the flow."""
-        return self.async_create_entry(
-            title=TITLE,
-            data={
-                CONF_URL: self.ws_address,
-                CONF_USB_PATH: self.usb_path,
-                CONF_NETWORK_KEY: self.network_key,
-                CONF_USE_ADDON: self.use_addon,
-                CONF_INTEGRATION_CREATED_ADDON: self.integration_created_addon,
-            },
-        )
-
-    async def async_step_on_supervisor(
-        self, user_input: dict[str, Any] | None = None
-    ) -> FlowResult:
-        """Handle logic when on Supervisor host."""
-        if user_input is None:
-            return self.async_show_form(
-                step_id="on_supervisor", data_schema=ON_SUPERVISOR_SCHEMA
-            )
-        if not user_input[CONF_USE_ADDON]:
-            return await self.async_step_manual()
-
-        self.use_addon = True
-
-        addon_info = await self._async_get_addon_info()
-
-        if addon_info.state == AddonState.RUNNING:
-            addon_config = addon_info.options
-            self.usb_path = addon_config[CONF_ADDON_DEVICE]
-            self.network_key = addon_config.get(CONF_ADDON_NETWORK_KEY, "")
-            return await self.async_step_finish_addon_setup()
-
-        if addon_info.state == AddonState.NOT_RUNNING:
-            return await self.async_step_configure_addon()
-
-        return await self.async_step_install_addon()
 
     async def async_step_install_addon(
         self, user_input: dict[str, Any] | None = None
@@ -213,10 +113,12 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         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")
 
         self.integration_created_addon = True
+        self.install_task = None
 
         return self.async_show_progress_done(next_step_id="configure_addon")
 
@@ -226,43 +128,6 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         """Add-on installation failed."""
         return self.async_abort(reason="addon_install_failed")
 
-    async def async_step_configure_addon(
-        self, user_input: dict[str, Any] | None = None
-    ) -> FlowResult:
-        """Ask for config for Z-Wave JS add-on."""
-        addon_info = await self._async_get_addon_info()
-        addon_config = addon_info.options
-
-        errors: dict[str, str] = {}
-
-        if user_input is not None:
-            self.network_key = user_input[CONF_NETWORK_KEY]
-            self.usb_path = user_input[CONF_USB_PATH]
-
-            new_addon_config = {
-                CONF_ADDON_DEVICE: self.usb_path,
-                CONF_ADDON_NETWORK_KEY: self.network_key,
-            }
-
-            if new_addon_config != addon_config:
-                await self._async_set_addon_config(new_addon_config)
-
-            return await self.async_step_start_addon()
-
-        usb_path = addon_config.get(CONF_ADDON_DEVICE, self.usb_path or "")
-        network_key = addon_config.get(CONF_ADDON_NETWORK_KEY, self.network_key or "")
-
-        data_schema = vol.Schema(
-            {
-                vol.Required(CONF_USB_PATH, default=usb_path): str,
-                vol.Optional(CONF_NETWORK_KEY, default=network_key): str,
-            }
-        )
-
-        return self.async_show_form(
-            step_id="configure_addon", data_schema=data_schema, errors=errors
-        )
-
     async def async_step_start_addon(
         self, user_input: dict[str, Any] | None = None
     ) -> FlowResult:
@@ -275,10 +140,12 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
         try:
             await self.start_task
-        except (CannotConnect, AddonError) as err:
+        except (CannotConnect, AddonError, AbortFlow) as err:
+            self.start_task = None
             _LOGGER.error(err)
             return self.async_show_progress_done(next_step_id="start_failed")
 
+        self.start_task = None
         return self.async_show_progress_done(next_step_id="finish_addon_setup")
 
     async def async_step_start_failed(
@@ -290,6 +157,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     async def _async_start_addon(self) -> None:
         """Start the Z-Wave JS add-on."""
         addon_manager: AddonManager = get_addon_manager(self.hass)
+        self.version_info = None
         try:
             await addon_manager.async_schedule_start_addon()
             # Sleep some seconds to let the add-on start properly before connecting.
@@ -301,7 +169,9 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                         self.ws_address = (
                             f"ws://{discovery_info['host']}:{discovery_info['port']}"
                         )
-                    await async_get_version_info(self.hass, self.ws_address)
+                    self.version_info = await async_get_version_info(
+                        self.hass, self.ws_address
+                    )
                 except (AbortFlow, CannotConnect) as err:
                     _LOGGER.debug(
                         "Add-on not ready yet, waiting %s seconds: %s",
@@ -315,9 +185,16 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         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)
+                self.flow_manager.async_configure(flow_id=self.flow_id)
             )
 
+    @abstractmethod
+    async def async_step_configure_addon(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Ask for config for Z-Wave JS add-on."""
+
+    @abstractmethod
     async def async_step_finish_addon_setup(
         self, user_input: dict[str, Any] | None = None
     ) -> FlowResult:
@@ -326,27 +203,6 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         Get add-on discovery info and server version info.
         Set unique id and abort if already configured.
         """
-        if not self.ws_address:
-            discovery_info = await self._async_get_addon_discovery_info()
-            self.ws_address = f"ws://{discovery_info['host']}:{discovery_info['port']}"
-
-        if not self.unique_id:
-            try:
-                version_info = await async_get_version_info(self.hass, self.ws_address)
-            except CannotConnect as err:
-                raise AbortFlow("cannot_connect") from err
-            await self.async_set_unique_id(
-                version_info.home_id, raise_on_progress=False
-            )
-
-        self._abort_if_unique_id_configured(
-            updates={
-                CONF_URL: self.ws_address,
-                CONF_USB_PATH: self.usb_path,
-                CONF_NETWORK_KEY: self.network_key,
-            }
-        )
-        return self._async_create_entry_from_vars()
 
     async def _async_get_addon_info(self) -> AddonInfo:
         """Return and cache Z-Wave JS add-on info."""
@@ -376,7 +232,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         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)
+                self.flow_manager.async_configure(flow_id=self.flow_id)
             )
 
     async def _async_get_addon_discovery_info(self) -> dict:
@@ -391,6 +247,204 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         return discovery_info_config
 
 
+class ConfigFlow(BaseZwaveJSFlow, config_entries.ConfigFlow, domain=DOMAIN):
+    """Handle a config flow for Z-Wave JS."""
+
+    VERSION = 1
+
+    def __init__(self) -> None:
+        """Set up flow instance."""
+        super().__init__()
+        self.use_addon = False
+
+    @property
+    def flow_manager(self) -> config_entries.ConfigEntriesFlowManager:
+        """Return the correct flow manager."""
+        return self.hass.config_entries.flow
+
+    async def async_step_user(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Handle the initial step."""
+        if is_hassio(self.hass):
+            return await self.async_step_on_supervisor()
+
+        return await self.async_step_manual()
+
+    async def async_step_manual(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Handle a manual configuration."""
+        if user_input is None:
+            return self.async_show_form(
+                step_id="manual", data_schema=get_manual_schema({})
+            )
+
+        errors = {}
+
+        try:
+            version_info = await validate_input(self.hass, user_input)
+        except InvalidInput as err:
+            errors["base"] = err.error
+        except Exception:  # pylint: disable=broad-except
+            _LOGGER.exception("Unexpected exception")
+            errors["base"] = "unknown"
+        else:
+            await self.async_set_unique_id(
+                version_info.home_id, raise_on_progress=False
+            )
+            # Make sure we disable any add-on handling
+            # if the controller is reconfigured in a manual step.
+            self._abort_if_unique_id_configured(
+                updates={
+                    **user_input,
+                    CONF_USE_ADDON: False,
+                    CONF_INTEGRATION_CREATED_ADDON: False,
+                }
+            )
+            self.ws_address = user_input[CONF_URL]
+            return self._async_create_entry_from_vars()
+
+        return self.async_show_form(
+            step_id="manual", data_schema=get_manual_schema(user_input), errors=errors
+        )
+
+    async def async_step_hassio(self, discovery_info: dict[str, Any]) -> FlowResult:
+        """Receive configuration from add-on discovery info.
+
+        This flow is triggered by the Z-Wave JS add-on.
+        """
+        self.ws_address = f"ws://{discovery_info['host']}:{discovery_info['port']}"
+        try:
+            version_info = await async_get_version_info(self.hass, self.ws_address)
+        except CannotConnect:
+            return self.async_abort(reason="cannot_connect")
+
+        await self.async_set_unique_id(version_info.home_id)
+        self._abort_if_unique_id_configured(updates={CONF_URL: self.ws_address})
+
+        return await self.async_step_hassio_confirm()
+
+    async def async_step_hassio_confirm(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Confirm the add-on discovery."""
+        if user_input is not None:
+            return await self.async_step_on_supervisor(
+                user_input={CONF_USE_ADDON: True}
+            )
+
+        return self.async_show_form(step_id="hassio_confirm")
+
+    async def async_step_on_supervisor(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Handle logic when on Supervisor host."""
+        if user_input is None:
+            return self.async_show_form(
+                step_id="on_supervisor", data_schema=ON_SUPERVISOR_SCHEMA
+            )
+        if not user_input[CONF_USE_ADDON]:
+            return await self.async_step_manual()
+
+        self.use_addon = True
+
+        addon_info = await self._async_get_addon_info()
+
+        if addon_info.state == AddonState.RUNNING:
+            addon_config = addon_info.options
+            self.usb_path = addon_config[CONF_ADDON_DEVICE]
+            self.network_key = addon_config.get(CONF_ADDON_NETWORK_KEY, "")
+            return await self.async_step_finish_addon_setup()
+
+        if addon_info.state == AddonState.NOT_RUNNING:
+            return await self.async_step_configure_addon()
+
+        return await self.async_step_install_addon()
+
+    async def async_step_configure_addon(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Ask for config for Z-Wave JS add-on."""
+        addon_info = await self._async_get_addon_info()
+        addon_config = addon_info.options
+
+        if user_input is not None:
+            self.network_key = user_input[CONF_NETWORK_KEY]
+            self.usb_path = user_input[CONF_USB_PATH]
+
+            new_addon_config = {
+                **addon_config,
+                CONF_ADDON_DEVICE: self.usb_path,
+                CONF_ADDON_NETWORK_KEY: self.network_key,
+            }
+
+            if new_addon_config != addon_config:
+                await self._async_set_addon_config(new_addon_config)
+
+            return await self.async_step_start_addon()
+
+        usb_path = addon_config.get(CONF_ADDON_DEVICE, self.usb_path or "")
+        network_key = addon_config.get(CONF_ADDON_NETWORK_KEY, self.network_key or "")
+
+        data_schema = vol.Schema(
+            {
+                vol.Required(CONF_USB_PATH, default=usb_path): str,
+                vol.Optional(CONF_NETWORK_KEY, default=network_key): str,
+            }
+        )
+
+        return self.async_show_form(step_id="configure_addon", data_schema=data_schema)
+
+    async def async_step_finish_addon_setup(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Prepare info needed to complete the config entry.
+
+        Get add-on discovery info and server version info.
+        Set unique id and abort if already configured.
+        """
+        if not self.ws_address:
+            discovery_info = await self._async_get_addon_discovery_info()
+            self.ws_address = f"ws://{discovery_info['host']}:{discovery_info['port']}"
+
+        if not self.unique_id:
+            if not self.version_info:
+                try:
+                    self.version_info = await async_get_version_info(
+                        self.hass, self.ws_address
+                    )
+                except CannotConnect as err:
+                    raise AbortFlow("cannot_connect") from err
+
+            await self.async_set_unique_id(
+                self.version_info.home_id, raise_on_progress=False
+            )
+
+        self._abort_if_unique_id_configured(
+            updates={
+                CONF_URL: self.ws_address,
+                CONF_USB_PATH: self.usb_path,
+                CONF_NETWORK_KEY: self.network_key,
+            }
+        )
+        return self._async_create_entry_from_vars()
+
+    @callback
+    def _async_create_entry_from_vars(self) -> FlowResult:
+        """Return a config entry for the flow."""
+        return self.async_create_entry(
+            title=TITLE,
+            data={
+                CONF_URL: self.ws_address,
+                CONF_USB_PATH: self.usb_path,
+                CONF_NETWORK_KEY: self.network_key,
+                CONF_USE_ADDON: self.use_addon,
+                CONF_INTEGRATION_CREATED_ADDON: self.integration_created_addon,
+            },
+        )
+
+
 class CannotConnect(exceptions.HomeAssistantError):
     """Indicate connection error."""