Update Hydrawise from the legacy API to the new GraphQL API (#106904)

* Update Hydrawise from the legacy API to the new GraphQL API.

* Cleanup

homeassistant/components/hydrawise/__init__.py

@@ -1,10 +1,11 @@
 """Support for Hydrawise cloud."""
 
-from pydrawise import legacy
+from pydrawise import auth, client
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_API_KEY, Platform
+from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryAuthFailed
 
 from .const import DOMAIN, SCAN_INTERVAL
 from .coordinator import HydrawiseDataUpdateCoordinator
@@ -14,8 +15,15 @@ PLATFORMS: list[Platform] = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.S
 
 async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Set up Hydrawise from a config entry."""
-    access_token = config_entry.data[CONF_API_KEY]
-    hydrawise = legacy.LegacyHydrawiseAsync(access_token)
+    if CONF_USERNAME not in config_entry.data or CONF_PASSWORD not in config_entry.data:
+        # The GraphQL API requires username and password to authenticate. If either is
+        # missing, reauth is required.
+        raise ConfigEntryAuthFailed
+
+    hydrawise = client.Hydrawise(
+        auth.Auth(config_entry.data[CONF_USERNAME], config_entry.data[CONF_PASSWORD])
+    )
+
     coordinator = HydrawiseDataUpdateCoordinator(hass, hydrawise, SCAN_INTERVAL)
     await coordinator.async_config_entry_first_refresh()
     hass.data.setdefault(DOMAIN, {})[config_entry.entry_id] = coordinator

homeassistant/components/hydrawise/config_flow.py

@@ -2,15 +2,16 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from typing import Any
 
 from aiohttp import ClientError
-from pydrawise import legacy
+from pydrawise import auth, client
+from pydrawise.exceptions import NotAuthorizedError
 import voluptuous as vol
 
-from homeassistant.config_entries import ConfigFlow, ConfigFlowResult
-from homeassistant.const import CONF_API_KEY
+from homeassistant.config_entries import ConfigEntry, ConfigFlow, ConfigFlowResult
+from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 
 from .const import DOMAIN, LOGGER
 
@@ -20,14 +21,26 @@ class HydrawiseConfigFlow(ConfigFlow, domain=DOMAIN):
 
     VERSION = 1
 
-    async def _create_entry(
-        self, api_key: str, *, on_failure: Callable[[str], ConfigFlowResult]
+    def __init__(self) -> None:
+        """Construct a ConfigFlow."""
+        self.reauth_entry: ConfigEntry | None = None
+
+    async def _create_or_update_entry(
+        self,
+        username: str,
+        password: str,
+        *,
+        on_failure: Callable[[str], ConfigFlowResult],
     ) -> ConfigFlowResult:
         """Create the config entry."""
-        api = legacy.LegacyHydrawiseAsync(api_key)
+
+        # Verify that the provided credentials work."""
+        api = client.Hydrawise(auth.Auth(username, password))
         try:
             # Skip fetching zones to save on metered API calls.
-            user = await api.get_user(fetch_zones=False)
+            user = await api.get_user()
+        except NotAuthorizedError:
+            return on_failure("invalid_auth")
         except TimeoutError:
             return on_failure("timeout_connect")
         except ClientError as ex:
@@ -35,17 +48,33 @@ class HydrawiseConfigFlow(ConfigFlow, domain=DOMAIN):
             return on_failure("cannot_connect")
 
         await self.async_set_unique_id(f"hydrawise-{user.customer_id}")
-        self._abort_if_unique_id_configured()
 
-        return self.async_create_entry(title="Hydrawise", data={CONF_API_KEY: api_key})
+        if not self.reauth_entry:
+            self._abort_if_unique_id_configured()
+            return self.async_create_entry(
+                title="Hydrawise",
+                data={CONF_USERNAME: username, CONF_PASSWORD: password},
+            )
+
+        self.hass.config_entries.async_update_entry(
+            self.reauth_entry,
+            data=self.reauth_entry.data
+            | {CONF_USERNAME: username, CONF_PASSWORD: password},
+        )
+        await self.hass.config_entries.async_reload(self.reauth_entry.entry_id)
+        return self.async_abort(reason="reauth_successful")
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Handle the initial setup."""
         if user_input is not None:
-            api_key = user_input[CONF_API_KEY]
-            return await self._create_entry(api_key, on_failure=self._show_form)
+            username = user_input[CONF_USERNAME]
+            password = user_input[CONF_PASSWORD]
+
+            return await self._create_or_update_entry(
+                username=username, password=password, on_failure=self._show_form
+            )
         return self._show_form()
 
     def _show_form(self, error_type: str | None = None) -> ConfigFlowResult:
@@ -54,6 +83,17 @@ class HydrawiseConfigFlow(ConfigFlow, domain=DOMAIN):
             errors["base"] = error_type
         return self.async_show_form(
             step_id="user",
-            data_schema=vol.Schema({vol.Required(CONF_API_KEY): str}),
+            data_schema=vol.Schema(
+                {vol.Required(CONF_USERNAME): str, vol.Required(CONF_PASSWORD): str}
+            ),
             errors=errors,
         )
+
+    async def async_step_reauth(
+        self, user_input: Mapping[str, Any]
+    ) -> ConfigFlowResult:
+        """Perform reauth after updating config to username/password."""
+        self.reauth_entry = self.hass.config_entries.async_get_entry(
+            self.context["entry_id"]
+        )
+        return await self.async_step_user()

homeassistant/components/hydrawise/strings.json

@@ -2,8 +2,11 @@
   "config": {
     "step": {
       "user": {
+        "title": "Hydrawise Login",
+        "description": "Please provide the username and password for your Hydrawise cloud account:",
         "data": {
-          "api_key": "[%key:common::config_flow::data::api_key%]"
+          "username": "[%key:common::config_flow::data::username%]",
+          "password": "[%key:common::config_flow::data::password%]"
         }
       }
     },
@@ -13,7 +16,8 @@
       "unknown": "[%key:common::config_flow::error::unknown%]"
     },
     "abort": {
-      "already_configured": "[%key:common::config_flow::abort::already_configured_service%]"
+      "already_configured": "[%key:common::config_flow::abort::already_configured_service%]",
+      "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]"
     }
   },
   "entity": {

tests/components/hydrawise/conftest.py

@@ -15,7 +15,7 @@ from pydrawise.schema import (
 import pytest
 
 from homeassistant.components.hydrawise.const import DOMAIN
-from homeassistant.const import CONF_API_KEY
+from homeassistant.const import CONF_API_KEY, CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 from homeassistant.util import dt as dt_util
 
@@ -32,7 +32,7 @@ def mock_setup_entry() -> Generator[AsyncMock, None, None]:
 
 
 @pytest.fixture
-def mock_pydrawise(
+def mock_legacy_pydrawise(
     user: User,
     controller: Controller,
     zones: list[Zone],
@@ -47,10 +47,32 @@ def mock_pydrawise(
         yield mock_pydrawise.return_value
 
 
+@pytest.fixture
+def mock_pydrawise(
+    mock_auth: AsyncMock,
+    user: User,
+    controller: Controller,
+    zones: list[Zone],
+) -> Generator[AsyncMock, None, None]:
+    """Mock Hydrawise."""
+    with patch("pydrawise.client.Hydrawise", autospec=True) as mock_pydrawise:
+        user.controllers = [controller]
+        controller.zones = zones
+        mock_pydrawise.return_value.get_user.return_value = user
+        yield mock_pydrawise.return_value
+
+
+@pytest.fixture
+def mock_auth() -> Generator[AsyncMock, None, None]:
+    """Mock pydrawise Auth."""
+    with patch("pydrawise.auth.Auth", autospec=True) as mock_auth:
+        yield mock_auth.return_value
+
+
 @pytest.fixture
 def user() -> User:
     """Hydrawise User fixture."""
-    return User(customer_id=12345)
+    return User(customer_id=12345, email="asdf@asdf.com")
 
 
 @pytest.fixture
@@ -102,7 +124,7 @@ def zones() -> list[Zone]:
 
 
 @pytest.fixture
-def mock_config_entry() -> MockConfigEntry:
+def mock_config_entry_legacy() -> MockConfigEntry:
     """Mock ConfigEntry."""
     return MockConfigEntry(
         title="Hydrawise",
@@ -111,6 +133,23 @@ def mock_config_entry() -> MockConfigEntry:
             CONF_API_KEY: "abc123",
         },
         unique_id="hydrawise-customerid",
+        version=1,
+    )
+
+
+@pytest.fixture
+def mock_config_entry() -> MockConfigEntry:
+    """Mock ConfigEntry."""
+    return MockConfigEntry(
+        title="Hydrawise",
+        domain=DOMAIN,
+        data={
+            CONF_USERNAME: "asfd@asdf.com",
+            CONF_PASSWORD: "__password__",
+        },
+        unique_id="hydrawise-customerid",
+        version=1,
+        minor_version=2,
     )
 
 

tests/components/hydrawise/test_config_flow.py

@@ -3,14 +3,18 @@
 from unittest.mock import AsyncMock
 
 from aiohttp import ClientError
+from pydrawise.exceptions import NotAuthorizedError
 from pydrawise.schema import User
 import pytest
 
 from homeassistant import config_entries
 from homeassistant.components.hydrawise.const import DOMAIN
+from homeassistant.const import CONF_API_KEY, CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
+from tests.common import MockConfigEntry
+
 pytestmark = pytest.mark.usefixtures("mock_setup_entry")
 
 
@@ -29,16 +33,20 @@ async def test_form(
     assert result["errors"] == {}
 
     result2 = await hass.config_entries.flow.async_configure(
-        result["flow_id"], {"api_key": "abc123"}
+        result["flow_id"],
+        {CONF_USERNAME: "asdf@asdf.com", CONF_PASSWORD: "__password__"},
     )
     mock_pydrawise.get_user.return_value = user
     await hass.async_block_till_done()
 
     assert result2["type"] is FlowResultType.CREATE_ENTRY
     assert result2["title"] == "Hydrawise"
-    assert result2["data"] == {"api_key": "abc123"}
+    assert result2["data"] == {
+        CONF_USERNAME: "asdf@asdf.com",
+        CONF_PASSWORD: "__password__",
+    }
     assert len(mock_setup_entry.mock_calls) == 1
-    mock_pydrawise.get_user.assert_called_once_with(fetch_zones=False)
+    mock_pydrawise.get_user.assert_called_once_with()
 
 
 async def test_form_api_error(
@@ -50,7 +58,7 @@ async def test_form_api_error(
     init_result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
-    data = {"api_key": "abc123"}
+    data = {CONF_USERNAME: "asdf@asdf.com", CONF_PASSWORD: "__password__"}
     result = await hass.config_entries.flow.async_configure(
         init_result["flow_id"], data
     )
@@ -71,7 +79,7 @@ async def test_form_connect_timeout(
     init_result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
-    data = {"api_key": "abc123"}
+    data = {CONF_USERNAME: "asdf@asdf.com", CONF_PASSWORD: "__password__"}
     result = await hass.config_entries.flow.async_configure(
         init_result["flow_id"], data
     )
@@ -83,3 +91,60 @@ async def test_form_connect_timeout(
     mock_pydrawise.get_user.return_value = user
     result2 = await hass.config_entries.flow.async_configure(result["flow_id"], data)
     assert result2["type"] is FlowResultType.CREATE_ENTRY
+
+
+async def test_form_not_authorized_error(
+    hass: HomeAssistant, mock_pydrawise: AsyncMock, user: User
+) -> None:
+    """Test we handle API errors."""
+    mock_pydrawise.get_user.side_effect = NotAuthorizedError
+
+    init_result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    data = {CONF_USERNAME: "asdf@asdf.com", 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": "invalid_auth"}
+
+    mock_pydrawise.get_user.reset_mock(side_effect=True)
+    mock_pydrawise.get_user.return_value = user
+    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_pydrawise: AsyncMock,
+) -> None:
+    """Test that re-authorization works."""
+    mock_config_entry = MockConfigEntry(
+        title="Hydrawise",
+        domain=DOMAIN,
+        data={
+            CONF_API_KEY: "__api_key__",
+        },
+        unique_id="hydrawise-12345",
+    )
+    mock_config_entry.add_to_hass(hass)
+
+    mock_config_entry.async_start_reauth(hass)
+    await hass.async_block_till_done()
+
+    flows = hass.config_entries.flow.async_progress()
+    assert len(flows) == 1
+    [result] = flows
+    assert result["step_id"] == "user"
+
+    result2 = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {CONF_USERNAME: "asdf@asdf.com", CONF_PASSWORD: "__password__"},
+    )
+    mock_pydrawise.get_user.return_value = user
+    await hass.async_block_till_done()
+
+    assert result2["type"] is FlowResultType.ABORT
+    assert result2["reason"] == "reauth_successful"

tests/components/hydrawise/test_init.py

@@ -19,3 +19,16 @@ async def test_connect_retry(
     await hass.config_entries.async_setup(mock_config_entry.entry_id)
     await hass.async_block_till_done()
     assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_update_version(
+    hass: HomeAssistant, mock_config_entry_legacy: MockConfigEntry
+) -> None:
+    """Test updating to the GaphQL API works."""
+    mock_config_entry_legacy.add_to_hass(hass)
+    await hass.config_entries.async_setup(mock_config_entry_legacy.entry_id)
+    await hass.async_block_till_done()
+    assert mock_config_entry_legacy.state is ConfigEntryState.SETUP_ERROR
+
+    # Make sure reauth flow has been initiated
+    assert any(mock_config_entry_legacy.async_get_active_flows(hass, {"reauth"}))