Catch Hydrawise authorization errors in the correct place (#132727)

homeassistant/components/hydrawise/config_flow.py

@@ -6,7 +6,7 @@
 from typing import Any
 
 from aiohttp import ClientError
-from pydrawise import auth, client
+from pydrawise import auth as pydrawise_auth, client
 from pydrawise.exceptions import NotAuthorizedError
 import voluptuous as vol
 
@@ -29,16 +29,21 @@ async def _create_or_update_entry(
         on_failure: Callable[[str], ConfigFlowResult],
     ) -> ConfigFlowResult:
         """Create the config entry."""
-
         # Verify that the provided credentials work."""
-        api = client.Hydrawise(auth.Auth(username, password))
+        auth = pydrawise_auth.Auth(username, password)
         try:
-            # Don't fetch zones because we don't need them yet.
-            user = await api.get_user(fetch_zones=False)
+            await auth.token()
         except NotAuthorizedError:
             return on_failure("invalid_auth")
         except TimeoutError:
             return on_failure("timeout_connect")
+
+        try:
+            api = client.Hydrawise(auth)
+            # Don't fetch zones because we don't need them yet.
+            user = await api.get_user(fetch_zones=False)
+        except TimeoutError:
+            return on_failure("timeout_connect")
         except ClientError as ex:
             LOGGER.error("Unable to connect to Hydrawise cloud service: %s", ex)
             return on_failure("cannot_connect")

tests/components/hydrawise/conftest.py

@@ -56,7 +56,6 @@ def mock_legacy_pydrawise(
 
 @pytest.fixture
 def mock_pydrawise(
-    mock_auth: AsyncMock,
     user: User,
     controller: Controller,
     zones: list[Zone],

tests/components/hydrawise/test_config_flow.py

@@ -21,6 +21,7 @@
 async def test_form(
     hass: HomeAssistant,
     mock_setup_entry: AsyncMock,
+    mock_auth: AsyncMock,
     mock_pydrawise: AsyncMock,
     user: User,
 ) -> None:
@@ -46,11 +47,12 @@ async def test_form(
         CONF_PASSWORD: "__password__",
     }
     assert len(mock_setup_entry.mock_calls) == 1
-    mock_pydrawise.get_user.assert_called_once_with(fetch_zones=False)
+    mock_auth.token.assert_awaited_once_with()
+    mock_pydrawise.get_user.assert_awaited_once_with(fetch_zones=False)
 
 
 async def test_form_api_error(
-    hass: HomeAssistant, mock_pydrawise: AsyncMock, user: User
+    hass: HomeAssistant, mock_auth: AsyncMock, mock_pydrawise: AsyncMock, user: User
 ) -> None:
     """Test we handle API errors."""
     mock_pydrawise.get_user.side_effect = ClientError("XXX")
@@ -71,8 +73,29 @@ async def test_form_api_error(
     assert result2["type"] is FlowResultType.CREATE_ENTRY
 
 
-async def test_form_connect_timeout(
-    hass: HomeAssistant, mock_pydrawise: AsyncMock, user: User
+async def test_form_auth_connect_timeout(
+    hass: HomeAssistant, mock_auth: AsyncMock, mock_pydrawise: AsyncMock
+) -> None:
+    """Test we handle API errors."""
+    mock_auth.token.side_effect = TimeoutError
+    init_result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    data = {CONF_USERNAME: "[email protected]", CONF_PASSWORD: "__password__"}
+    result = await hass.config_entries.flow.async_configure(
+        init_result["flow_id"], data
+    )
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["errors"] == {"base": "timeout_connect"}
+
+    mock_auth.token.reset_mock(side_effect=True)
+    result2 = await hass.config_entries.flow.async_configure(result["flow_id"], data)
+    assert result2["type"] is FlowResultType.CREATE_ENTRY
+
+
+async def test_form_client_connect_timeout(
+    hass: HomeAssistant, mock_auth: AsyncMock, mock_pydrawise: AsyncMock, user: User
 ) -> None:
     """Test we handle API errors."""
     mock_pydrawise.get_user.side_effect = TimeoutError
@@ -94,10 +117,10 @@ async def test_form_connect_timeout(
 
 
 async def test_form_not_authorized_error(
-    hass: HomeAssistant, mock_pydrawise: AsyncMock, user: User
+    hass: HomeAssistant, mock_auth: AsyncMock, mock_pydrawise: AsyncMock
 ) -> None:
     """Test we handle API errors."""
-    mock_pydrawise.get_user.side_effect = NotAuthorizedError
+    mock_auth.token.side_effect = NotAuthorizedError
 
     init_result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
@@ -109,15 +132,15 @@ async def test_form_not_authorized_error(
     assert result["type"] is FlowResultType.FORM
     assert result["errors"] == {"base": "invalid_auth"}
 
-    mock_pydrawise.get_user.reset_mock(side_effect=True)
-    mock_pydrawise.get_user.return_value = user
+    mock_auth.token.reset_mock(side_effect=True)
     result2 = await hass.config_entries.flow.async_configure(result["flow_id"], data)
     assert result2["type"] is FlowResultType.CREATE_ENTRY
 
 
 async def test_reauth(
     hass: HomeAssistant,
     user: User,
+    mock_auth: AsyncMock,
     mock_pydrawise: AsyncMock,
 ) -> None:
     """Test that re-authorization works."""