Improve zwave_js config flow (#46906)

homeassistant/components/zwave_js/config_flow.py

@@ -8,8 +8,18 @@ from async_timeout import timeout
 import voluptuous as vol
 from zwave_js_server.version import VersionInfo, get_server_version
 
-from homeassistant import config_entries, core, exceptions
+from homeassistant import config_entries, exceptions
+from homeassistant.components.hassio import (
+    async_get_addon_discovery_info,
+    async_get_addon_info,
+    async_install_addon,
+    async_set_addon_options,
+    async_start_addon,
+    is_hassio,
+)
+from homeassistant.components.hassio.handler import HassioAPIError
 from homeassistant.const import CONF_URL
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.data_entry_flow import AbortFlow
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
@@ -29,13 +39,14 @@ CONF_USB_PATH = "usb_path"
 DEFAULT_URL = "ws://localhost:3000"
 TITLE = "Z-Wave JS"
 
-ADDON_SETUP_TIME = 10
+ADDON_SETUP_TIMEOUT = 5
+ADDON_SETUP_TIMEOUT_ROUNDS = 4
 
 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})
 
 
-async def validate_input(hass: core.HomeAssistant, user_input: dict) -> VersionInfo:
+async def validate_input(hass: HomeAssistant, user_input: dict) -> VersionInfo:
     """Validate if the user input allows us to connect."""
     ws_address = user_input[CONF_URL]
 
@@ -48,9 +59,7 @@ async def validate_input(hass: core.HomeAssistant, user_input: dict) -> VersionI
         raise InvalidInput("cannot_connect") from err
 
 
-async def async_get_version_info(
+async def async_get_version_info(hass: HomeAssistant, ws_address: str) -> VersionInfo:
-    hass: core.HomeAssistant, ws_address: str
-) -> VersionInfo:
     """Return Z-Wave JS version info."""
     async with timeout(10):
         try:
@@ -58,7 +67,9 @@ async def async_get_version_info(
                 ws_address, async_get_clientsession(hass)
             )
         except (asyncio.TimeoutError, aiohttp.ClientError) as err:
-            _LOGGER.error("Failed to connect to Z-Wave JS server: %s", err)
+            # We don't want to spam the log if the add-on isn't started
+            # or takes a long time to start.
+            _LOGGER.debug("Failed to connect to Z-Wave JS server: %s", err)
             raise CannotConnect from err
 
     return version_info
@@ -72,7 +83,6 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
     def __init__(self) -> None:
         """Set up flow instance."""
-        self.addon_config: Optional[dict] = None
         self.network_key: Optional[str] = None
         self.usb_path: Optional[str] = None
         self.use_addon = False
@@ -80,12 +90,14 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         # If we install the add-on we should uninstall it on entry remove.
         self.integration_created_addon = False
         self.install_task: Optional[asyncio.Task] = None
+        self.start_task: Optional[asyncio.Task] = None
 
     async def async_step_user(
         self, user_input: Optional[Dict[str, Any]] = None
     ) -> Dict[str, Any]:
         """Handle the initial step."""
-        if self.hass.components.hassio.is_hassio():
+        assert self.hass  # typing
+        if is_hassio(self.hass):  # type: ignore  # no-untyped-call
             return await self.async_step_on_supervisor()
 
         return await self.async_step_manual()
@@ -120,7 +132,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             step_id="manual", data_schema=STEP_USER_DATA_SCHEMA, errors=errors
         )
 
-    async def async_step_hassio(  # type: ignore
+    async def async_step_hassio(  # type: ignore # override
         self, discovery_info: Dict[str, Any]
     ) -> Dict[str, Any]:
         """Receive configuration from add-on discovery info.
@@ -149,6 +161,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
         return self.async_show_form(step_id="hassio_confirm")
 
+    @callback
     def _async_create_entry_from_vars(self) -> Dict[str, Any]:
         """Return a config entry for the flow."""
         return self.async_create_entry(
@@ -176,28 +189,13 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         self.use_addon = True
 
         if await self._async_is_addon_running():
-            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:
-                    return self.async_abort(reason="cannot_connect")
-                await self.async_set_unique_id(
-                    version_info.home_id, raise_on_progress=False
-                )
-
-            self._abort_if_unique_id_configured()
             addon_config = await self._async_get_addon_config()
             self.usb_path = addon_config[CONF_ADDON_DEVICE]
             self.network_key = addon_config.get(CONF_ADDON_NETWORK_KEY, "")
-            return self._async_create_entry_from_vars()
+            return await self.async_step_finish_addon_setup()
 
         if await self._async_is_addon_installed():
-            return await self.async_step_start_addon()
+            return await self.async_step_configure_addon()
 
         return await self.async_step_install_addon()
 
@@ -213,13 +211,13 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
         try:
             await self.install_task
-        except self.hass.components.hassio.HassioAPIError as err:
+        except HassioAPIError as err:
             _LOGGER.error("Failed to install Z-Wave JS add-on: %s", err)
             return self.async_show_progress_done(next_step_id="install_failed")
 
         self.integration_created_addon = True
 
-        return self.async_show_progress_done(next_step_id="start_addon")
+        return self.async_show_progress_done(next_step_id="configure_addon")
 
     async def async_step_install_failed(
         self, user_input: Optional[Dict[str, Any]] = None
@@ -227,14 +225,13 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         """Add-on installation failed."""
         return self.async_abort(reason="addon_install_failed")
 
-    async def async_step_start_addon(
+    async def async_step_configure_addon(
         self, user_input: Optional[Dict[str, Any]] = None
     ) -> Dict[str, Any]:
-        """Ask for config and start Z-Wave JS add-on."""
+        """Ask for config for Z-Wave JS add-on."""
-        if self.addon_config is None:
+        addon_config = await self._async_get_addon_config()
-            self.addon_config = await self._async_get_addon_config()
 
-        errors = {}
+        errors: Dict[str, str] = {}
 
         if user_input is not None:
             self.network_key = user_input[CONF_NETWORK_KEY]
@@ -245,40 +242,13 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 CONF_ADDON_NETWORK_KEY: self.network_key,
             }
 
-            if new_addon_config != self.addon_config:
+            if new_addon_config != addon_config:
                 await self._async_set_addon_config(new_addon_config)
 
-            try:
+            return await self.async_step_start_addon()
-                await self.hass.components.hassio.async_start_addon(ADDON_SLUG)
-            except self.hass.components.hassio.HassioAPIError as err:
-                _LOGGER.error("Failed to start Z-Wave JS add-on: %s", err)
-                errors["base"] = "addon_start_failed"
-            else:
-                # Sleep some seconds to let the add-on start properly before connecting.
-                await asyncio.sleep(ADDON_SETUP_TIME)
-                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:
+        usb_path = addon_config.get(CONF_ADDON_DEVICE, self.usb_path or "")
-                    try:
+        network_key = addon_config.get(CONF_ADDON_NETWORK_KEY, self.network_key or "")
-                        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, raise_on_progress=False
-                    )
-
-                self._abort_if_unique_id_configured()
-                return self._async_create_entry_from_vars()
-
-        usb_path = self.addon_config.get(CONF_ADDON_DEVICE, self.usb_path or "")
-        network_key = self.addon_config.get(
-            CONF_ADDON_NETWORK_KEY, self.network_key or ""
-        )
 
         data_schema = vol.Schema(
             {
@@ -288,16 +258,95 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         )
 
         return self.async_show_form(
-            step_id="start_addon", data_schema=data_schema, errors=errors
+            step_id="configure_addon", data_schema=data_schema, errors=errors
         )
 
+    async def async_step_start_addon(
+        self, user_input: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """Start Z-Wave JS add-on."""
+        assert self.hass
+        if not self.start_task:
+            self.start_task = self.hass.async_create_task(self._async_start_addon())
+            return self.async_show_progress(
+                step_id="start_addon", progress_action="start_addon"
+            )
+
+        try:
+            await self.start_task
+        except (CannotConnect, HassioAPIError) as err:
+            _LOGGER.error("Failed to start Z-Wave JS add-on: %s", err)
+            return self.async_show_progress_done(next_step_id="start_failed")
+
+        return self.async_show_progress_done(next_step_id="finish_addon_setup")
+
+    async def async_step_start_failed(
+        self, user_input: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """Add-on start failed."""
+        return self.async_abort(reason="addon_start_failed")
+
+    async def _async_start_addon(self) -> None:
+        """Start the Z-Wave JS add-on."""
+        assert self.hass
+        try:
+            await async_start_addon(self.hass, ADDON_SLUG)
+            # 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 self.ws_address:
+                        discovery_info = await self._async_get_addon_discovery_info()
+                        self.ws_address = (
+                            f"ws://{discovery_info['host']}:{discovery_info['port']}"
+                        )
+                    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",
+                        ADDON_SETUP_TIMEOUT,
+                        err,
+                    )
+                else:
+                    break
+            else:
+                raise CannotConnect("Failed to start 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)
+            )
+
+    async def async_step_finish_addon_setup(
+        self, user_input: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """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.
+        """
+        assert self.hass
+        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()
+        return self._async_create_entry_from_vars()
+
     async def _async_get_addon_info(self) -> dict:
         """Return and cache Z-Wave JS add-on info."""
         try:
-            addon_info: dict = await self.hass.components.hassio.async_get_addon_info(
+            addon_info: dict = await async_get_addon_info(self.hass, ADDON_SLUG)
-                ADDON_SLUG
+        except HassioAPIError as err:
-            )
-        except self.hass.components.hassio.HassioAPIError as err:
             _LOGGER.error("Failed to get Z-Wave JS add-on info: %s", err)
             raise AbortFlow("addon_info_failed") from err
 
@@ -322,17 +371,15 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         """Set Z-Wave JS add-on config."""
         options = {"options": config}
         try:
-            await self.hass.components.hassio.async_set_addon_options(
+            await async_set_addon_options(self.hass, ADDON_SLUG, options)
-                ADDON_SLUG, options
+        except HassioAPIError as err:
-            )
-        except self.hass.components.hassio.HassioAPIError as err:
             _LOGGER.error("Failed to set Z-Wave JS add-on config: %s", err)
             raise AbortFlow("addon_set_config_failed") from err
 
     async def _async_install_addon(self) -> None:
         """Install the Z-Wave JS add-on."""
         try:
-            await self.hass.components.hassio.async_install_addon(ADDON_SLUG)
+            await async_install_addon(self.hass, ADDON_SLUG)
         finally:
             # Continue the flow after show progress when the task is done.
             self.hass.async_create_task(
@@ -342,12 +389,8 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     async def _async_get_addon_discovery_info(self) -> dict:
         """Return add-on discovery info."""
         try:
-            discovery_info: dict = (
+            discovery_info = await async_get_addon_discovery_info(self.hass, ADDON_SLUG)
-                await self.hass.components.hassio.async_get_addon_discovery_info(
+        except HassioAPIError as err:
-                    ADDON_SLUG
-                )
-            )
-        except self.hass.components.hassio.HassioAPIError as err:
             _LOGGER.error("Failed to get Z-Wave JS add-on discovery info: %s", err)
             raise AbortFlow("addon_get_discovery_info_failed") from err
 

homeassistant/components/zwave_js/strings.json

@@ -15,13 +15,14 @@
       "install_addon": {
         "title": "The Z-Wave JS add-on installation has started"
       },
-      "start_addon": {
+      "configure_addon": {
         "title": "Enter the Z-Wave JS add-on configuration",
         "data": {
           "usb_path": "[%key:common::config_flow::data::usb_path%]",
           "network_key": "Network Key"
         }
       },
+      "start_addon": { "title": "The Z-Wave JS add-on is starting." },
       "hassio_confirm": {
         "title": "Set up Z-Wave JS integration with the Z-Wave JS add-on"
       }
@@ -38,12 +39,14 @@
       "addon_info_failed": "Failed to get Z-Wave JS add-on info.",
       "addon_install_failed": "Failed to install the Z-Wave JS add-on.",
       "addon_set_config_failed": "Failed to set Z-Wave JS configuration.",
+      "addon_start_failed": "Failed to start the Z-Wave JS add-on.",
       "addon_get_discovery_info_failed": "Failed to get Z-Wave JS add-on discovery info.",
       "addon_missing_discovery_info": "Missing Z-Wave JS add-on discovery info.",
       "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]"
     },
     "progress": {
-      "install_addon": "Please wait while the Z-Wave JS add-on installation finishes. This can take several minutes."
+      "install_addon": "Please wait while the Z-Wave JS add-on installation finishes. This can take several minutes.",
+      "start_addon": "Please wait while the Z-Wave JS add-on start completes. This may take some seconds."
     }
   }
 }

homeassistant/components/zwave_js/translations/en.json

@@ -6,6 +6,7 @@
             "addon_install_failed": "Failed to install the Z-Wave JS add-on.",
             "addon_missing_discovery_info": "Missing Z-Wave JS add-on discovery info.",
             "addon_set_config_failed": "Failed to set Z-Wave JS configuration.",
+            "addon_start_failed": "Failed to start the Z-Wave JS add-on.",
             "already_configured": "Device is already configured",
             "already_in_progress": "Configuration flow is already in progress",
             "cannot_connect": "Failed to connect"
@@ -17,9 +18,17 @@
             "unknown": "Unexpected error"
         },
         "progress": {
-            "install_addon": "Please wait while the Z-Wave JS add-on installation finishes. This can take several minutes."
+            "install_addon": "Please wait while the Z-Wave JS add-on installation finishes. This can take several minutes.",
+            "start_addon": "Please wait while the Z-Wave JS add-on start completes. This may take some seconds."
         },
         "step": {
+            "configure_addon": {
+                "data": {
+                    "network_key": "Network Key",
+                    "usb_path": "USB Device Path"
+                },
+                "title": "Enter the Z-Wave JS add-on configuration"
+            },
             "hassio_confirm": {
                 "title": "Set up Z-Wave JS integration with the Z-Wave JS add-on"
             },
@@ -39,18 +48,9 @@
                 "title": "Select connection method"
             },
             "start_addon": {
-                "data": {
+                "title": "The Z-Wave JS add-on is starting."
-                    "network_key": "Network Key",
-                    "usb_path": "USB Device Path"
-                },
-                "title": "Enter the Z-Wave JS add-on configuration"
-            },
-            "user": {
-                "data": {
-                    "url": "URL"
-                }
             }
         }
     },
     "title": "Z-Wave JS"
-}
+}

tests/components/zwave_js/conftest.py

@@ -1,7 +1,7 @@
 """Provide common Z-Wave JS fixtures."""
 import asyncio
 import json
-from unittest.mock import DEFAULT, AsyncMock, patch
+from unittest.mock import AsyncMock, patch
 
 import pytest
 from zwave_js_server.event import Event
@@ -22,29 +22,6 @@ async def device_registry_fixture(hass):
     return await async_get_device_registry(hass)
 
 
-@pytest.fixture(name="discovery_info")
-def discovery_info_fixture():
-    """Return the discovery info from the supervisor."""
-    return DEFAULT
-
-
-@pytest.fixture(name="discovery_info_side_effect")
-def discovery_info_side_effect_fixture():
-    """Return the discovery info from the supervisor."""
-    return None
-
-
-@pytest.fixture(name="get_addon_discovery_info")
-def mock_get_addon_discovery_info(discovery_info, discovery_info_side_effect):
-    """Mock get add-on discovery info."""
-    with patch(
-        "homeassistant.components.hassio.async_get_addon_discovery_info",
-        side_effect=discovery_info_side_effect,
-        return_value=discovery_info,
-    ) as get_addon_discovery_info:
-        yield get_addon_discovery_info
-
-
 @pytest.fixture(name="controller_state", scope="session")
 def controller_state_fixture():
     """Load the controller state fixture data."""

tests/components/zwave_js/test_config_flow.py

@@ -1,6 +1,6 @@
 """Test the Z-Wave JS config flow."""
 import asyncio
-from unittest.mock import patch
+from unittest.mock import DEFAULT, patch
 
 import pytest
 from zwave_js_server.version import VersionInfo
@@ -22,10 +22,35 @@ ADDON_DISCOVERY_INFO = {
 @pytest.fixture(name="supervisor")
 def mock_supervisor_fixture():
     """Mock Supervisor."""
-    with patch("homeassistant.components.hassio.is_hassio", return_value=True):
+    with patch(
+        "homeassistant.components.zwave_js.config_flow.is_hassio", return_value=True
+    ):
         yield
 
 
+@pytest.fixture(name="discovery_info")
+def discovery_info_fixture():
+    """Return the discovery info from the supervisor."""
+    return DEFAULT
+
+
+@pytest.fixture(name="discovery_info_side_effect")
+def discovery_info_side_effect_fixture():
+    """Return the discovery info from the supervisor."""
+    return None
+
+
+@pytest.fixture(name="get_addon_discovery_info")
+def mock_get_addon_discovery_info(discovery_info, discovery_info_side_effect):
+    """Mock get add-on discovery info."""
+    with patch(
+        "homeassistant.components.zwave_js.config_flow.async_get_addon_discovery_info",
+        side_effect=discovery_info_side_effect,
+        return_value=discovery_info,
+    ) as get_addon_discovery_info:
+        yield get_addon_discovery_info
+
+
 @pytest.fixture(name="addon_info_side_effect")
 def addon_info_side_effect_fixture():
     """Return the add-on info side effect."""
@@ -36,7 +61,7 @@ def addon_info_side_effect_fixture():
 def mock_addon_info(addon_info_side_effect):
     """Mock Supervisor add-on info."""
     with patch(
-        "homeassistant.components.hassio.async_get_addon_info",
+        "homeassistant.components.zwave_js.config_flow.async_get_addon_info",
         side_effect=addon_info_side_effect,
     ) as addon_info:
         addon_info.return_value = {}
@@ -75,7 +100,7 @@ def set_addon_options_side_effect_fixture():
 def mock_set_addon_options(set_addon_options_side_effect):
     """Mock set add-on options."""
     with patch(
-        "homeassistant.components.hassio.async_set_addon_options",
+        "homeassistant.components.zwave_js.config_flow.async_set_addon_options",
         side_effect=set_addon_options_side_effect,
     ) as set_options:
         yield set_options
@@ -84,7 +109,9 @@ def mock_set_addon_options(set_addon_options_side_effect):
 @pytest.fixture(name="install_addon")
 def mock_install_addon():
     """Mock install add-on."""
-    with patch("homeassistant.components.hassio.async_install_addon") as install_addon:
+    with patch(
+        "homeassistant.components.zwave_js.config_flow.async_install_addon"
+    ) as install_addon:
         yield install_addon
 
 
@@ -98,7 +125,7 @@ def start_addon_side_effect_fixture():
 def mock_start_addon(start_addon_side_effect):
     """Mock start add-on."""
     with patch(
-        "homeassistant.components.hassio.async_start_addon",
+        "homeassistant.components.zwave_js.config_flow.async_start_addon",
         side_effect=start_addon_side_effect,
     ) as start_addon:
         yield start_addon
@@ -130,7 +157,7 @@ def mock_get_server_version(server_version_side_effect):
 def mock_addon_setup_time():
     """Mock add-on setup sleep time."""
     with patch(
-        "homeassistant.components.zwave_js.config_flow.ADDON_SETUP_TIME", new=0
+        "homeassistant.components.zwave_js.config_flow.ADDON_SETUP_TIMEOUT", new=0
     ) as addon_setup_time:
         yield addon_setup_time
 
@@ -399,12 +426,47 @@ async def test_discovery_addon_not_running(
 
     result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
 
-    assert result["step_id"] == "start_addon"
     assert result["type"] == "form"
+    assert result["step_id"] == "configure_addon"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"usb_path": "/test", "network_key": "abc123"}
+    )
+
+    assert result["type"] == "progress"
+    assert result["step_id"] == "start_addon"
+
+    with patch(
+        "homeassistant.components.zwave_js.async_setup", return_value=True
+    ) as mock_setup, patch(
+        "homeassistant.components.zwave_js.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry:
+        await hass.async_block_till_done()
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
+        await hass.async_block_till_done()
+
+    assert result["type"] == "create_entry"
+    assert result["title"] == TITLE
+    assert result["data"] == {
+        "url": "ws://host1:3001",
+        "usb_path": "/test",
+        "network_key": "abc123",
+        "use_addon": True,
+        "integration_created_addon": False,
+    }
+    assert len(mock_setup.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 1
 
 
 async def test_discovery_addon_not_installed(
-    hass, supervisor, addon_installed, install_addon, addon_options
+    hass,
+    supervisor,
+    addon_installed,
+    install_addon,
+    addon_options,
+    set_addon_options,
+    start_addon,
 ):
     """Test discovery with add-on not installed."""
     addon_installed.return_value["version"] = None
@@ -429,8 +491,37 @@ async def test_discovery_addon_not_installed(
     result = await hass.config_entries.flow.async_configure(result["flow_id"])
 
     assert result["type"] == "form"
+    assert result["step_id"] == "configure_addon"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"usb_path": "/test", "network_key": "abc123"}
+    )
+
+    assert result["type"] == "progress"
     assert result["step_id"] == "start_addon"
 
+    with patch(
+        "homeassistant.components.zwave_js.async_setup", return_value=True
+    ) as mock_setup, patch(
+        "homeassistant.components.zwave_js.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry:
+        await hass.async_block_till_done()
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
+        await hass.async_block_till_done()
+
+    assert result["type"] == "create_entry"
+    assert result["title"] == TITLE
+    assert result["data"] == {
+        "url": "ws://host1:3001",
+        "usb_path": "/test",
+        "network_key": "abc123",
+        "use_addon": True,
+        "integration_created_addon": True,
+    }
+    assert len(mock_setup.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 1
+
 
 async def test_not_addon(hass, supervisor):
     """Test opting out of add-on on Supervisor."""
@@ -559,10 +650,13 @@ async def test_addon_running_failures(
     hass,
     supervisor,
     addon_running,
+    addon_options,
     get_addon_discovery_info,
     abort_reason,
 ):
     """Test all failures when add-on is running."""
+    addon_options["device"] = "/test"
+    addon_options["network_key"] = "abc123"
     await setup.async_setup_component(hass, "persistent_notification", {})
 
     result = await hass.config_entries.flow.async_init(
@@ -582,9 +676,11 @@ async def test_addon_running_failures(
 
 @pytest.mark.parametrize("discovery_info", [{"config": ADDON_DISCOVERY_INFO}])
 async def test_addon_running_already_configured(
-    hass, supervisor, addon_running, get_addon_discovery_info
+    hass, supervisor, addon_running, addon_options, get_addon_discovery_info
 ):
     """Test that only one unique instance is allowed when add-on is running."""
+    addon_options["device"] = "/test"
+    addon_options["network_key"] = "abc123"
     entry = MockConfigEntry(domain=DOMAIN, data={}, title=TITLE, unique_id=1234)
     entry.add_to_hass(hass)
 
@@ -629,6 +725,13 @@ async def test_addon_installed(
     )
 
     assert result["type"] == "form"
+    assert result["step_id"] == "configure_addon"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"usb_path": "/test", "network_key": "abc123"}
+    )
+
+    assert result["type"] == "progress"
     assert result["step_id"] == "start_addon"
 
     with patch(
@@ -637,9 +740,8 @@ async def test_addon_installed(
         "homeassistant.components.zwave_js.async_setup_entry",
         return_value=True,
     ) as mock_setup_entry:
-        result = await hass.config_entries.flow.async_configure(
+        await hass.async_block_till_done()
-            result["flow_id"], {"usb_path": "/test", "network_key": "abc123"}
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
-        )
         await hass.async_block_till_done()
 
     assert result["type"] == "create_entry"
@@ -683,40 +785,32 @@ async def test_addon_installed_start_failure(
     )
 
     assert result["type"] == "form"
-    assert result["step_id"] == "start_addon"
+    assert result["step_id"] == "configure_addon"
 
     result = await hass.config_entries.flow.async_configure(
         result["flow_id"], {"usb_path": "/test", "network_key": "abc123"}
     )
 
-    assert result["type"] == "form"
+    assert result["type"] == "progress"
-    assert result["errors"] == {"base": "addon_start_failed"}
+    assert result["step_id"] == "start_addon"
+
+    await hass.async_block_till_done()
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    assert result["type"] == "abort"
+    assert result["reason"] == "addon_start_failed"
 
 
 @pytest.mark.parametrize(
-    "set_addon_options_side_effect, start_addon_side_effect, discovery_info, "
+    "discovery_info, server_version_side_effect",
-    "server_version_side_effect, abort_reason",
     [
         (
-            HassioAPIError(),
-            None,
-            {"config": ADDON_DISCOVERY_INFO},
-            None,
-            "addon_set_config_failed",
-        ),
-        (
-            None,
-            None,
             {"config": ADDON_DISCOVERY_INFO},
             asyncio.TimeoutError,
-            "cannot_connect",
         ),
         (
             None,
             None,
-            None,
-            None,
-            "addon_missing_discovery_info",
         ),
     ],
 )
@@ -728,7 +822,6 @@ async def test_addon_installed_failures(
     set_addon_options,
     start_addon,
     get_addon_discovery_info,
-    abort_reason,
 ):
     """Test all failures when add-on is installed."""
     await setup.async_setup_component(hass, "persistent_notification", {})
@@ -745,14 +838,58 @@ async def test_addon_installed_failures(
     )
 
     assert result["type"] == "form"
+    assert result["step_id"] == "configure_addon"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"usb_path": "/test", "network_key": "abc123"}
+    )
+
+    assert result["type"] == "progress"
     assert result["step_id"] == "start_addon"
 
+    await hass.async_block_till_done()
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    assert result["type"] == "abort"
+    assert result["reason"] == "addon_start_failed"
+
+
+@pytest.mark.parametrize(
+    "set_addon_options_side_effect, discovery_info",
+    [(HassioAPIError(), {"config": ADDON_DISCOVERY_INFO})],
+)
+async def test_addon_installed_set_options_failure(
+    hass,
+    supervisor,
+    addon_installed,
+    addon_options,
+    set_addon_options,
+    start_addon,
+    get_addon_discovery_info,
+):
+    """Test all failures when add-on is installed."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+
+    assert result["type"] == "form"
+    assert result["step_id"] == "on_supervisor"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"use_addon": True}
+    )
+
+    assert result["type"] == "form"
+    assert result["step_id"] == "configure_addon"
+
     result = await hass.config_entries.flow.async_configure(
         result["flow_id"], {"usb_path": "/test", "network_key": "abc123"}
     )
 
     assert result["type"] == "abort"
-    assert result["reason"] == abort_reason
+    assert result["reason"] == "addon_set_config_failed"
 
 
 @pytest.mark.parametrize("discovery_info", [{"config": ADDON_DISCOVERY_INFO}])
@@ -782,12 +919,18 @@ async def test_addon_installed_already_configured(
     )
 
     assert result["type"] == "form"
-    assert result["step_id"] == "start_addon"
+    assert result["step_id"] == "configure_addon"
 
     result = await hass.config_entries.flow.async_configure(
         result["flow_id"], {"usb_path": "/test", "network_key": "abc123"}
     )
 
+    assert result["type"] == "progress"
+    assert result["step_id"] == "start_addon"
+
+    await hass.async_block_till_done()
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
     assert result["type"] == "abort"
     assert result["reason"] == "already_configured"
 
@@ -819,6 +962,7 @@ async def test_addon_not_installed(
     )
 
     assert result["type"] == "progress"
+    assert result["step_id"] == "install_addon"
 
     # Make sure the flow continues when the progress task is done.
     await hass.async_block_till_done()
@@ -826,6 +970,13 @@ async def test_addon_not_installed(
     result = await hass.config_entries.flow.async_configure(result["flow_id"])
 
     assert result["type"] == "form"
+    assert result["step_id"] == "configure_addon"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"usb_path": "/test", "network_key": "abc123"}
+    )
+
+    assert result["type"] == "progress"
     assert result["step_id"] == "start_addon"
 
     with patch(
@@ -834,9 +985,8 @@ async def test_addon_not_installed(
         "homeassistant.components.zwave_js.async_setup_entry",
         return_value=True,
     ) as mock_setup_entry:
-        result = await hass.config_entries.flow.async_configure(
+        await hass.async_block_till_done()
-            result["flow_id"], {"usb_path": "/test", "network_key": "abc123"}
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
-        )
         await hass.async_block_till_done()
 
     assert result["type"] == "create_entry"