From db8f597f1052dda14f5415423bcea571141047ea Mon Sep 17 00:00:00 2001
From: Martin Hjelmare <marhje52@gmail.com>
Date: Wed, 24 Feb 2021 16:22:58 +0100
Subject: [PATCH] Fix zwave_js config flow server version timeout (#46990)

---
 .../components/zwave_js/config_flow.py        | 15 ++++++-----
 tests/components/zwave_js/test_config_flow.py | 27 +++++++++++++++----
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/homeassistant/components/zwave_js/config_flow.py b/homeassistant/components/zwave_js/config_flow.py
index fe79796edf1..cc19fb85d3a 100644
--- a/homeassistant/components/zwave_js/config_flow.py
+++ b/homeassistant/components/zwave_js/config_flow.py
@@ -41,6 +41,7 @@ TITLE = "Z-Wave JS"
 
 ADDON_SETUP_TIMEOUT = 5
 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})
@@ -61,16 +62,16 @@ async def validate_input(hass: HomeAssistant, user_input: dict) -> VersionInfo:
 
 async def async_get_version_info(hass: HomeAssistant, ws_address: str) -> VersionInfo:
     """Return Z-Wave JS version info."""
-    async with timeout(10):
-        try:
+    try:
+        async with timeout(SERVER_VERSION_TIMEOUT):
             version_info: VersionInfo = await get_server_version(
                 ws_address, async_get_clientsession(hass)
             )
-        except (asyncio.TimeoutError, aiohttp.ClientError) as 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
+    except (asyncio.TimeoutError, aiohttp.ClientError) as 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
 
diff --git a/tests/components/zwave_js/test_config_flow.py b/tests/components/zwave_js/test_config_flow.py
index 3c956d42a27..73057f3fe21 100644
--- a/tests/components/zwave_js/test_config_flow.py
+++ b/tests/components/zwave_js/test_config_flow.py
@@ -7,7 +7,7 @@ from zwave_js_server.version import VersionInfo
 
 from homeassistant import config_entries, setup
 from homeassistant.components.hassio.handler import HassioAPIError
-from homeassistant.components.zwave_js.config_flow import TITLE
+from homeassistant.components.zwave_js.config_flow import SERVER_VERSION_TIMEOUT, TITLE
 from homeassistant.components.zwave_js.const import DOMAIN
 
 from tests.common import MockConfigEntry
@@ -138,7 +138,7 @@ def server_version_side_effect_fixture():
 
 
 @pytest.fixture(name="get_server_version", autouse=True)
-def mock_get_server_version(server_version_side_effect):
+def mock_get_server_version(server_version_side_effect, server_version_timeout):
     """Mock server version."""
     version_info = VersionInfo(
         driver_version="mock-driver-version",
@@ -149,10 +149,19 @@ def mock_get_server_version(server_version_side_effect):
         "homeassistant.components.zwave_js.config_flow.get_server_version",
         side_effect=server_version_side_effect,
         return_value=version_info,
-    ) as mock_version:
+    ) as mock_version, patch(
+        "homeassistant.components.zwave_js.config_flow.SERVER_VERSION_TIMEOUT",
+        new=server_version_timeout,
+    ):
         yield mock_version
 
 
+@pytest.fixture(name="server_version_timeout")
+def mock_server_version_timeout():
+    """Patch the timeout for getting server version."""
+    return SERVER_VERSION_TIMEOUT
+
+
 @pytest.fixture(name="addon_setup_time", autouse=True)
 def mock_addon_setup_time():
     """Mock add-on setup sleep time."""
@@ -198,22 +207,30 @@ async def test_manual(hass):
     assert result2["result"].unique_id == 1234
 
 
+async def slow_server_version(*args):
+    """Simulate a slow server version."""
+    await asyncio.sleep(0.1)
+
+
 @pytest.mark.parametrize(
-    "url, server_version_side_effect, error",
+    "url, server_version_side_effect, server_version_timeout, error",
     [
         (
             "not-ws-url",
             None,
+            SERVER_VERSION_TIMEOUT,
             "invalid_ws_url",
         ),
         (
             "ws://localhost:3000",
-            asyncio.TimeoutError,
+            slow_server_version,
+            0,
             "cannot_connect",
         ),
         (
             "ws://localhost:3000",
             Exception("Boom"),
+            SERVER_VERSION_TIMEOUT,
             "unknown",
         ),
     ],