Migrate Hydrawise to an async client library (#103636)

* Migrate Hydrawise to an async client library

* Changes requested during review

* Additional changes requested during review

homeassistant/components/hydrawise/__init__.py

@@ -51,7 +51,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
 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.LegacyHydrawise(access_token, load_on_init=False)
+    hydrawise = legacy.LegacyHydrawiseAsync(access_token)
     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/binary_sensor.py

@@ -1,7 +1,7 @@
 """Support for Hydrawise sprinkler binary sensors."""
 from __future__ import annotations
 
-from pydrawise.legacy import LegacyHydrawise
+from pydrawise.schema import Zone
 import voluptuous as vol
 
 from homeassistant.components.binary_sensor import (
@@ -69,26 +69,16 @@ async def async_setup_entry(
     coordinator: HydrawiseDataUpdateCoordinator = hass.data[DOMAIN][
         config_entry.entry_id
     ]
-    hydrawise: LegacyHydrawise = coordinator.api
+    entities = []
-
+    for controller in coordinator.data.controllers:
-    entities = [
+        entities.append(
-        HydrawiseBinarySensor(
+            HydrawiseBinarySensor(coordinator, BINARY_SENSOR_STATUS, controller)
-            data=hydrawise.current_controller,
-            coordinator=coordinator,
-            description=BINARY_SENSOR_STATUS,
-            device_id_key="controller_id",
         )
-    ]
+        for zone in controller.zones:
-
-    # create a sensor for each zone
-    for zone in hydrawise.relays:
             for description in BINARY_SENSOR_TYPES:
                 entities.append(
-                HydrawiseBinarySensor(
+                    HydrawiseBinarySensor(coordinator, description, controller, zone)
-                    data=zone, coordinator=coordinator, description=description
                 )
-            )
-
     async_add_entities(entities)
 
 
@@ -100,5 +90,5 @@ class HydrawiseBinarySensor(HydrawiseEntity, BinarySensorEntity):
         if self.entity_description.key == "status":
             self._attr_is_on = self.coordinator.last_update_success
         elif self.entity_description.key == "is_watering":
-            relay_data = self.coordinator.api.relays_by_zone_number[self.data["relay"]]
+            zone: Zone = self.zone
-            self._attr_is_on = relay_data["timestr"] == "Now"
+            self._attr_is_on = zone.scheduled_runs.current_run is not None

homeassistant/components/hydrawise/config_flow.py

@@ -5,8 +5,8 @@ from __future__ import annotations
 from collections.abc import Callable
 from typing import Any
 
+from aiohttp import ClientError
 from pydrawise import legacy
-from requests.exceptions import ConnectTimeout, HTTPError
 import voluptuous as vol
 
 from homeassistant import config_entries
@@ -27,20 +27,17 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         self, api_key: str, *, on_failure: Callable[[str], FlowResult]
     ) -> FlowResult:
         """Create the config entry."""
+        api = legacy.LegacyHydrawiseAsync(api_key)
         try:
-            api = await self.hass.async_add_executor_job(
+            # Skip fetching zones to save on metered API calls.
-                legacy.LegacyHydrawise, api_key
+            user = await api.get_user(fetch_zones=False)
-            )
+        except TimeoutError:
-        except ConnectTimeout:
             return on_failure("timeout_connect")
-        except HTTPError as ex:
+        except ClientError as ex:
             LOGGER.error("Unable to connect to Hydrawise cloud service: %s", ex)
             return on_failure("cannot_connect")
 
-        if not api.status:
+        await self.async_set_unique_id(f"hydrawise-{user.customer_id}")
-            return on_failure("unknown")
-
-        await self.async_set_unique_id(f"hydrawise-{api.customer_id}")
         self._abort_if_unique_id_configured()
 
         return self.async_create_entry(title="Hydrawise", data={CONF_API_KEY: api_key})

homeassistant/components/hydrawise/coordinator.py

@@ -4,26 +4,25 @@ from __future__ import annotations
 
 from datetime import timedelta
 
-from pydrawise.legacy import LegacyHydrawise
+from pydrawise import HydrawiseBase
+from pydrawise.schema import User
 
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .const import DOMAIN, LOGGER
 
 
-class HydrawiseDataUpdateCoordinator(DataUpdateCoordinator[None]):
+class HydrawiseDataUpdateCoordinator(DataUpdateCoordinator[User]):
     """The Hydrawise Data Update Coordinator."""
 
     def __init__(
-        self, hass: HomeAssistant, api: LegacyHydrawise, scan_interval: timedelta
+        self, hass: HomeAssistant, api: HydrawiseBase, scan_interval: timedelta
     ) -> None:
         """Initialize HydrawiseDataUpdateCoordinator."""
         super().__init__(hass, LOGGER, name=DOMAIN, update_interval=scan_interval)
         self.api = api
 
-    async def _async_update_data(self) -> None:
+    async def _async_update_data(self) -> User:
         """Fetch the latest data from Hydrawise."""
-        result = await self.hass.async_add_executor_job(self.api.update_controller_info)
+        return await self.api.get_user()
-        if not result:
-            raise UpdateFailed("Failed to refresh Hydrawise data")

homeassistant/components/hydrawise/entity.py

@@ -1,7 +1,7 @@
 """Base classes for Hydrawise entities."""
 from __future__ import annotations
 
-from typing import Any
+from pydrawise.schema import Controller, Zone
 
 from homeassistant.core import callback
 from homeassistant.helpers.device_registry import DeviceInfo
@@ -20,23 +20,25 @@ class HydrawiseEntity(CoordinatorEntity[HydrawiseDataUpdateCoordinator]):
 
     def __init__(
         self,
-        *,
-        data: dict[str, Any],
         coordinator: HydrawiseDataUpdateCoordinator,
         description: EntityDescription,
-        device_id_key: str = "relay_id",
+        controller: Controller,
+        zone: Zone | None = None,
     ) -> None:
         """Initialize the Hydrawise entity."""
         super().__init__(coordinator=coordinator)
-        self.data = data
         self.entity_description = description
-        self._device_id = str(data.get(device_id_key))
+        self.controller = controller
+        self.zone = zone
+        self._device_id = str(controller.id if zone is None else zone.id)
         self._attr_unique_id = f"{self._device_id}_{description.key}"
         self._attr_device_info = DeviceInfo(
             identifiers={(DOMAIN, self._device_id)},
-            name=data["name"],
+            name=controller.name if zone is None else zone.name,
             manufacturer=MANUFACTURER,
         )
+        if zone is not None:
+            self._attr_device_info["via_device"] = (DOMAIN, str(controller.id))
         self._update_attrs()
 
     def _update_attrs(self) -> None:

homeassistant/components/hydrawise/sensor.py

@@ -1,6 +1,9 @@
 """Support for Hydrawise sprinkler sensors."""
 from __future__ import annotations
 
+from datetime import datetime
+
+from pydrawise.schema import Zone
 import voluptuous as vol
 
 from homeassistant.components.sensor import (
@@ -71,27 +74,30 @@ async def async_setup_entry(
     coordinator: HydrawiseDataUpdateCoordinator = hass.data[DOMAIN][
         config_entry.entry_id
     ]
-    entities = [
+    async_add_entities(
-        HydrawiseSensor(data=zone, coordinator=coordinator, description=description)
+        HydrawiseSensor(coordinator, description, controller, zone)
-        for zone in coordinator.api.relays
+        for controller in coordinator.data.controllers
+        for zone in controller.zones
         for description in SENSOR_TYPES
-    ]
+    )
-    async_add_entities(entities)
 
 
 class HydrawiseSensor(HydrawiseEntity, SensorEntity):
     """A sensor implementation for Hydrawise device."""
 
+    zone: Zone
+
     def _update_attrs(self) -> None:
         """Update state attributes."""
-        relay_data = self.coordinator.api.relays_by_zone_number[self.data["relay"]]
         if self.entity_description.key == "watering_time":
-            if relay_data["timestr"] == "Now":
+            if (current_run := self.zone.scheduled_runs.current_run) is not None:
-                self._attr_native_value = int(relay_data["run"] / 60)
+                self._attr_native_value = int(
+                    current_run.remaining_time.total_seconds() / 60
+                )
             else:
                 self._attr_native_value = 0
-        else:  # _sensor_type == 'next_cycle'
+        elif self.entity_description.key == "next_cycle":
-            next_cycle = min(relay_data["time"], TWO_YEAR_SECONDS)
+            if (next_run := self.zone.scheduled_runs.next_run) is not None:
-            self._attr_native_value = dt_util.utc_from_timestamp(
+                self._attr_native_value = dt_util.as_utc(next_run.start_time)
-                dt_util.as_timestamp(dt_util.now()) + next_cycle
+            else:
-            )
+                self._attr_native_value = datetime.max.replace(tzinfo=dt_util.UTC)

homeassistant/components/hydrawise/switch.py

@@ -1,8 +1,10 @@
 """Support for Hydrawise cloud switches."""
 from __future__ import annotations
 
+from datetime import timedelta
 from typing import Any
 
+from pydrawise.schema import Zone
 import voluptuous as vol
 
 from homeassistant.components.switch import (
@@ -17,6 +19,7 @@ from homeassistant.core import HomeAssistant
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
+from homeassistant.util import dt as dt_util
 
 from .const import (
     ALLOWED_WATERING_TIME,
@@ -76,62 +79,44 @@ async def async_setup_entry(
     coordinator: HydrawiseDataUpdateCoordinator = hass.data[DOMAIN][
         config_entry.entry_id
     ]
-    default_watering_timer = DEFAULT_WATERING_TIME
+    async_add_entities(
-
+        HydrawiseSwitch(coordinator, description, controller, zone)
-    entities = [
+        for controller in coordinator.data.controllers
-        HydrawiseSwitch(
+        for zone in controller.zones
-            data=zone,
-            coordinator=coordinator,
-            description=description,
-            default_watering_timer=default_watering_timer,
-        )
-        for zone in coordinator.api.relays
         for description in SWITCH_TYPES
-    ]
+    )
-
-    async_add_entities(entities)
 
 
 class HydrawiseSwitch(HydrawiseEntity, SwitchEntity):
     """A switch implementation for Hydrawise device."""
 
-    def __init__(
+    zone: Zone
-        self,
-        *,
-        data: dict[str, Any],
-        coordinator: HydrawiseDataUpdateCoordinator,
-        description: SwitchEntityDescription,
-        default_watering_timer: int,
-    ) -> None:
-        """Initialize a switch for Hydrawise device."""
-        super().__init__(data=data, coordinator=coordinator, description=description)
-        self._default_watering_timer = default_watering_timer
 
-    def turn_on(self, **kwargs: Any) -> None:
+    async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn the device on."""
-        zone_number = self.data["relay"]
         if self.entity_description.key == "manual_watering":
-            self.coordinator.api.run_zone(self._default_watering_timer, zone_number)
+            await self.coordinator.api.start_zone(
+                self.zone, custom_run_duration=DEFAULT_WATERING_TIME
+            )
         elif self.entity_description.key == "auto_watering":
-            self.coordinator.api.suspend_zone(0, zone_number)
+            await self.coordinator.api.resume_zone(self.zone)
         self._attr_is_on = True
         self.async_write_ha_state()
 
-    def turn_off(self, **kwargs: Any) -> None:
+    async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn the device off."""
-        zone_number = self.data["relay"]
         if self.entity_description.key == "manual_watering":
-            self.coordinator.api.run_zone(0, zone_number)
+            await self.coordinator.api.stop_zone(self.zone)
         elif self.entity_description.key == "auto_watering":
-            self.coordinator.api.suspend_zone(365, zone_number)
+            await self.coordinator.api.suspend_zone(
+                self.zone, dt_util.now() + timedelta(days=365)
+            )
         self._attr_is_on = False
         self.async_write_ha_state()
 
     def _update_attrs(self) -> None:
         """Update state attributes."""
-        zone_number = self.data["relay"]
-        timestr = self.coordinator.api.relays_by_zone_number[zone_number]["timestr"]
         if self.entity_description.key == "manual_watering":
-            self._attr_is_on = timestr == "Now"
+            self._attr_is_on = self.zone.scheduled_runs.current_run is not None
         elif self.entity_description.key == "auto_watering":
-            self._attr_is_on = timestr not in {"", "Now"}
+            self._attr_is_on = self.zone.status.suspended_until is None

tests/components/hydrawise/conftest.py

@@ -1,14 +1,23 @@
 """Common fixtures for the Hydrawise tests."""
 
-from collections.abc import Generator
+from collections.abc import Awaitable, Callable, Generator
-from typing import Any
+from datetime import datetime, timedelta
-from unittest.mock import AsyncMock, Mock, patch
+from unittest.mock import AsyncMock, patch
 
+from pydrawise.schema import (
+    Controller,
+    ControllerHardware,
+    ScheduledZoneRun,
+    ScheduledZoneRuns,
+    User,
+    Zone,
+)
 import pytest
 
 from homeassistant.components.hydrawise.const import DOMAIN
 from homeassistant.const import CONF_API_KEY
 from homeassistant.core import HomeAssistant
+from homeassistant.util import dt as dt_util
 
 from tests.common import MockConfigEntry
 
@@ -24,59 +33,71 @@ def mock_setup_entry() -> Generator[AsyncMock, None, None]:
 
 @pytest.fixture
 def mock_pydrawise(
-    mock_controller: dict[str, Any],
+    user: User,
-    mock_zones: list[dict[str, Any]],
+    controller: Controller,
-) -> Generator[Mock, None, None]:
+    zones: list[Zone],
-    """Mock LegacyHydrawise."""
+) -> Generator[AsyncMock, None, None]:
-    with patch("pydrawise.legacy.LegacyHydrawise", autospec=True) as mock_pydrawise:
+    """Mock LegacyHydrawiseAsync."""
-        mock_pydrawise.return_value.controller_info = {"controllers": [mock_controller]}
+    with patch(
-        mock_pydrawise.return_value.current_controller = mock_controller
+        "pydrawise.legacy.LegacyHydrawiseAsync", autospec=True
-        mock_pydrawise.return_value.controller_status = {"relays": mock_zones}
+    ) as mock_pydrawise:
-        mock_pydrawise.return_value.relays = mock_zones
+        user.controllers = [controller]
-        mock_pydrawise.return_value.relays_by_zone_number = {
+        controller.zones = zones
-            r["relay"]: r for r in mock_zones
+        mock_pydrawise.return_value.get_user.return_value = user
-        }
         yield mock_pydrawise.return_value
 
 
 @pytest.fixture
-def mock_controller() -> dict[str, Any]:
+def user() -> User:
-    """Mock Hydrawise controller."""
+    """Hydrawise User fixture."""
-    return {
+    return User(customer_id=12345)
-        "name": "Home Controller",
-        "last_contact": 1693292420,
-        "serial_number": "0310b36090",
-        "controller_id": 52496,
-        "status": "Unknown",
-    }
 
 
 @pytest.fixture
-def mock_zones() -> list[dict[str, Any]]:
+def controller() -> Controller:
-    """Mock Hydrawise zones."""
+    """Hydrawise Controller fixture."""
+    return Controller(
+        id=52496,
+        name="Home Controller",
+        hardware=ControllerHardware(
+            serial_number="0310b36090",
+        ),
+        last_contact_time=datetime.fromtimestamp(1693292420),
+        online=True,
+    )
+
+
+@pytest.fixture
+def zones() -> list[Zone]:
+    """Hydrawise zone fixtures."""
     return [
-        {
+        Zone(
-            "name": "Zone One",
+            name="Zone One",
-            "period": 259200,
+            number=1,
-            "relay": 1,
+            id=5965394,
-            "relay_id": 5965394,
+            scheduled_runs=ScheduledZoneRuns(
-            "run": 1800,
+                summary="",
-            "stop": 1,
+                current_run=None,
-            "time": 330597,
+                next_run=ScheduledZoneRun(
-            "timestr": "Sat",
+                    start_time=dt_util.now() + timedelta(seconds=330597),
-            "type": 1,
+                    end_time=dt_util.now()
-        },
+                    + timedelta(seconds=330597)
-        {
+                    + timedelta(seconds=1800),
-            "name": "Zone Two",
+                    normal_duration=timedelta(seconds=1800),
-            "period": 259200,
+                    duration=timedelta(seconds=1800),
-            "relay": 2,
+                ),
-            "relay_id": 5965395,
+            ),
-            "run": 1788,
+        ),
-            "stop": 1,
+        Zone(
-            "time": 1,
+            name="Zone Two",
-            "timestr": "Now",
+            number=2,
-            "type": 106,
+            id=5965395,
-        },
+            scheduled_runs=ScheduledZoneRuns(
+                current_run=ScheduledZoneRun(
+                    remaining_time=timedelta(seconds=1788),
+                ),
+            ),
+        ),
     ]
 
 
@@ -95,13 +116,25 @@ def mock_config_entry() -> MockConfigEntry:
 
 @pytest.fixture
 async def mock_added_config_entry(
-    hass: HomeAssistant,
+    mock_add_config_entry: Callable[[], Awaitable[MockConfigEntry]]
-    mock_config_entry: MockConfigEntry,
-    mock_pydrawise: Mock,
 ) -> MockConfigEntry:
     """Mock ConfigEntry that's been added to HA."""
+    return await mock_add_config_entry()
+
+
+@pytest.fixture
+async def mock_add_config_entry(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_pydrawise: AsyncMock,
+) -> Callable[[], Awaitable[MockConfigEntry]]:
+    """Callable that creates a mock ConfigEntry that's been added to HA."""
+
+    async def callback() -> MockConfigEntry:
         mock_config_entry.add_to_hass(hass)
         await hass.config_entries.async_setup(mock_config_entry.entry_id)
         await hass.async_block_till_done()
         assert DOMAIN in hass.config_entries.async_domains()
         return mock_config_entry
+
+    return callback

tests/components/hydrawise/test_binary_sensor.py

@@ -1,8 +1,9 @@
 """Test Hydrawise binary_sensor."""
 
 from datetime import timedelta
-from unittest.mock import Mock
+from unittest.mock import AsyncMock
 
+from aiohttp import ClientError
 from freezegun.api import FrozenDateTimeFactory
 
 from homeassistant.components.hydrawise.const import SCAN_INTERVAL
@@ -33,12 +34,13 @@ async def test_states(
 async def test_update_data_fails(
     hass: HomeAssistant,
     mock_added_config_entry: MockConfigEntry,
-    mock_pydrawise: Mock,
+    mock_pydrawise: AsyncMock,
     freezer: FrozenDateTimeFactory,
 ) -> None:
     """Test that no data from the API sets the correct connectivity."""
     # Make the coordinator refresh data.
-    mock_pydrawise.update_controller_info.return_value = None
+    mock_pydrawise.get_user.reset_mock(return_value=True)
+    mock_pydrawise.get_user.side_effect = ClientError
     freezer.tick(SCAN_INTERVAL + timedelta(seconds=30))
     async_fire_time_changed(hass)
     await hass.async_block_till_done()

tests/components/hydrawise/test_config_flow.py

@@ -1,9 +1,10 @@
 """Test the Hydrawise config flow."""
 
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock
 
+from aiohttp import ClientError
+from pydrawise.schema import User
 import pytest
-from requests.exceptions import ConnectTimeout, HTTPError
 
 from homeassistant import config_entries
 from homeassistant.components.hydrawise.const import DOMAIN
@@ -17,9 +18,11 @@ from tests.common import MockConfigEntry
 pytestmark = pytest.mark.usefixtures("mock_setup_entry")
 
 
-@patch("pydrawise.legacy.LegacyHydrawise")
 async def test_form(
-    mock_api: MagicMock, hass: HomeAssistant, mock_setup_entry: AsyncMock
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_pydrawise: AsyncMock,
+    user: User,
 ) -> None:
     """Test we get the form."""
     result = await hass.config_entries.flow.async_init(
@@ -32,19 +35,22 @@ async def test_form(
     result2 = await hass.config_entries.flow.async_configure(
         result["flow_id"], {"api_key": "abc123"}
     )
-    mock_api.return_value.customer_id = 12345
+    mock_pydrawise.get_user.return_value = user
     await hass.async_block_till_done()
 
     assert result2["type"] == FlowResultType.CREATE_ENTRY
     assert result2["title"] == "Hydrawise"
     assert result2["data"] == {"api_key": "abc123"}
     assert len(mock_setup_entry.mock_calls) == 1
+    mock_pydrawise.get_user.assert_called_once_with(fetch_zones=False)
 
 
-@patch("pydrawise.legacy.LegacyHydrawise")
+async def test_form_api_error(
-async def test_form_api_error(mock_api: MagicMock, hass: HomeAssistant) -> None:
+    hass: HomeAssistant, mock_pydrawise: AsyncMock, user: User
+) -> None:
     """Test we handle API errors."""
-    mock_api.side_effect = HTTPError
+    mock_pydrawise.get_user.side_effect = ClientError("XXX")
+
     init_result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
@@ -55,15 +61,17 @@ async def test_form_api_error(mock_api: MagicMock, hass: HomeAssistant) -> None:
     assert result["type"] == FlowResultType.FORM
     assert result["errors"] == {"base": "cannot_connect"}
 
-    mock_api.side_effect = None
+    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"] == FlowResultType.CREATE_ENTRY
 
 
-@patch("pydrawise.legacy.LegacyHydrawise")
+async def test_form_connect_timeout(
-async def test_form_connect_timeout(mock_api: MagicMock, hass: HomeAssistant) -> None:
+    hass: HomeAssistant, mock_pydrawise: AsyncMock, user: User
+) -> None:
     """Test we handle API errors."""
-    mock_api.side_effect = ConnectTimeout
+    mock_pydrawise.get_user.side_effect = TimeoutError
     init_result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
@@ -75,15 +83,17 @@ async def test_form_connect_timeout(mock_api: MagicMock, hass: HomeAssistant) ->
     assert result["type"] == FlowResultType.FORM
     assert result["errors"] == {"base": "timeout_connect"}
 
-    mock_api.side_effect = None
+    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"] == FlowResultType.CREATE_ENTRY
 
 
-@patch("pydrawise.legacy.LegacyHydrawise")
+async def test_flow_import_success(
-async def test_flow_import_success(mock_api: MagicMock, hass: HomeAssistant) -> None:
+    hass: HomeAssistant, mock_pydrawise: AsyncMock, user: User
+) -> None:
     """Test that we can import a YAML config."""
-    mock_api.return_value.status = "All good!"
+    mock_pydrawise.get_user.return_value = User
     result = await hass.config_entries.flow.async_init(
         DOMAIN,
         context={"source": config_entries.SOURCE_IMPORT},
@@ -107,9 +117,11 @@ async def test_flow_import_success(mock_api: MagicMock, hass: HomeAssistant) ->
     assert issue.translation_key == "deprecated_yaml"
 
 
-@patch("pydrawise.legacy.LegacyHydrawise", side_effect=HTTPError)
+async def test_flow_import_api_error(
-async def test_flow_import_api_error(mock_api: MagicMock, hass: HomeAssistant) -> None:
+    hass: HomeAssistant, mock_pydrawise: AsyncMock
+) -> None:
     """Test that we handle API errors on YAML import."""
+    mock_pydrawise.get_user.side_effect = ClientError
     result = await hass.config_entries.flow.async_init(
         DOMAIN,
         context={"source": config_entries.SOURCE_IMPORT},
@@ -129,11 +141,11 @@ async def test_flow_import_api_error(mock_api: MagicMock, hass: HomeAssistant) -
     assert issue.translation_key == "deprecated_yaml_import_issue"
 
 
-@patch("pydrawise.legacy.LegacyHydrawise", side_effect=ConnectTimeout)
 async def test_flow_import_connect_timeout(
-    mock_api: MagicMock, hass: HomeAssistant
+    hass: HomeAssistant, mock_pydrawise: AsyncMock
 ) -> None:
     """Test that we handle connection timeouts on YAML import."""
+    mock_pydrawise.get_user.side_effect = TimeoutError
     result = await hass.config_entries.flow.async_init(
         DOMAIN,
         context={"source": config_entries.SOURCE_IMPORT},
@@ -153,32 +165,8 @@ async def test_flow_import_connect_timeout(
     assert issue.translation_key == "deprecated_yaml_import_issue"
 
 
-@patch("pydrawise.legacy.LegacyHydrawise")
-async def test_flow_import_no_status(mock_api: MagicMock, hass: HomeAssistant) -> None:
-    """Test we handle a lack of API status on YAML import."""
-    mock_api.return_value.status = None
-    result = await hass.config_entries.flow.async_init(
-        DOMAIN,
-        context={"source": config_entries.SOURCE_IMPORT},
-        data={
-            CONF_API_KEY: "__api_key__",
-            CONF_SCAN_INTERVAL: 120,
-        },
-    )
-    await hass.async_block_till_done()
-    assert result["type"] == FlowResultType.ABORT
-    assert result["reason"] == "unknown"
-
-    issue_registry = ir.async_get(hass)
-    issue = issue_registry.async_get_issue(
-        DOMAIN, "deprecated_yaml_import_issue_unknown"
-    )
-    assert issue.translation_key == "deprecated_yaml_import_issue"
-
-
-@patch("pydrawise.legacy.LegacyHydrawise")
 async def test_flow_import_already_imported(
-    mock_api: MagicMock, hass: HomeAssistant
+    hass: HomeAssistant, mock_pydrawise: AsyncMock, user: User
 ) -> None:
     """Test that we can handle a YAML config already imported."""
     mock_config_entry = MockConfigEntry(
@@ -187,12 +175,12 @@ async def test_flow_import_already_imported(
         data={
             CONF_API_KEY: "__api_key__",
         },
-        unique_id="hydrawise-CUSTOMER_ID",
+        unique_id="hydrawise-12345",
     )
     mock_config_entry.add_to_hass(hass)
 
-    mock_api.return_value.customer_id = "CUSTOMER_ID"
+    mock_pydrawise.get_user.return_value = user
-    mock_api.return_value.status = "All good!"
+
     result = await hass.config_entries.flow.async_init(
         DOMAIN,
         context={"source": config_entries.SOURCE_IMPORT},

tests/components/hydrawise/test_init.py

@@ -1,8 +1,8 @@
 """Tests for the Hydrawise integration."""
 
-from unittest.mock import Mock
+from unittest.mock import AsyncMock
 
-from requests.exceptions import HTTPError
+from aiohttp import ClientError
 
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import CONF_ACCESS_TOKEN
@@ -13,11 +13,10 @@ from homeassistant.setup import async_setup_component
 from tests.common import MockConfigEntry
 
 
-async def test_setup_import_success(hass: HomeAssistant, mock_pydrawise: Mock) -> None:
+async def test_setup_import_success(
+    hass: HomeAssistant, mock_pydrawise: AsyncMock
+) -> None:
     """Test that setup with a YAML config triggers an import and warning."""
-    mock_pydrawise.update_controller_info.return_value = True
-    mock_pydrawise.customer_id = 12345
-    mock_pydrawise.status = "unknown"
     config = {"hydrawise": {CONF_ACCESS_TOKEN: "_access-token_"}}
     assert await async_setup_component(hass, "hydrawise", config)
     await hass.async_block_till_done()
@@ -30,21 +29,10 @@ async def test_setup_import_success(hass: HomeAssistant, mock_pydrawise: Mock) -
 
 
 async def test_connect_retry(
-    hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_pydrawise: Mock
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_pydrawise: AsyncMock
 ) -> None:
     """Test that a connection error triggers a retry."""
-    mock_pydrawise.update_controller_info.side_effect = HTTPError
+    mock_pydrawise.get_user.side_effect = ClientError
-    mock_config_entry.add_to_hass(hass)
-    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_setup_no_data(
-    hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_pydrawise: Mock
-) -> None:
-    """Test that no data from the API triggers a retry."""
-    mock_pydrawise.update_controller_info.return_value = False
     mock_config_entry.add_to_hass(hass)
     await hass.config_entries.async_setup(mock_config_entry.entry_id)
     await hass.async_block_till_done()

tests/components/hydrawise/test_sensor.py

@@ -1,6 +1,9 @@
 """Test Hydrawise sensor."""
 
+from collections.abc import Awaitable, Callable
+
 from freezegun.api import FrozenDateTimeFactory
+from pydrawise.schema import Zone
 import pytest
 
 from homeassistant.core import HomeAssistant
@@ -26,3 +29,18 @@ async def test_states(
     next_cycle = hass.states.get("sensor.zone_one_next_cycle")
     assert next_cycle is not None
     assert next_cycle.state == "2023-10-04T19:49:57+00:00"
+
+
+@pytest.mark.freeze_time("2023-10-01 00:00:00+00:00")
+async def test_suspended_state(
+    hass: HomeAssistant,
+    zones: list[Zone],
+    mock_add_config_entry: Callable[[], Awaitable[MockConfigEntry]],
+) -> None:
+    """Test sensor states."""
+    zones[0].scheduled_runs.next_run = None
+    await mock_add_config_entry()
+
+    next_cycle = hass.states.get("sensor.zone_one_next_cycle")
+    assert next_cycle is not None
+    assert next_cycle.state == "9999-12-31T23:59:59+00:00"

tests/components/hydrawise/test_switch.py

@@ -1,12 +1,16 @@
 """Test Hydrawise switch."""
 
-from unittest.mock import Mock
+from datetime import timedelta
+from unittest.mock import AsyncMock
 
-from freezegun.api import FrozenDateTimeFactory
+from pydrawise.schema import Zone
+import pytest
 
+from homeassistant.components.hydrawise.const import DEFAULT_WATERING_TIME
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.const import ATTR_ENTITY_ID, SERVICE_TURN_OFF, SERVICE_TURN_ON
 from homeassistant.core import HomeAssistant
+from homeassistant.util import dt as dt_util
 
 from tests.common import MockConfigEntry
 
@@ -14,7 +18,6 @@ from tests.common import MockConfigEntry
 async def test_states(
     hass: HomeAssistant,
     mock_added_config_entry: MockConfigEntry,
-    freezer: FrozenDateTimeFactory,
 ) -> None:
     """Test switch states."""
     watering1 = hass.states.get("switch.zone_one_manual_watering")
@@ -31,11 +34,14 @@ async def test_states(
 
     auto_watering2 = hass.states.get("switch.zone_two_automatic_watering")
     assert auto_watering2 is not None
-    assert auto_watering2.state == "off"
+    assert auto_watering2.state == "on"
 
 
 async def test_manual_watering_services(
-    hass: HomeAssistant, mock_added_config_entry: MockConfigEntry, mock_pydrawise: Mock
+    hass: HomeAssistant,
+    mock_added_config_entry: MockConfigEntry,
+    mock_pydrawise: AsyncMock,
+    zones: list[Zone],
 ) -> None:
     """Test Manual Watering services."""
     await hass.services.async_call(
@@ -44,7 +50,9 @@ async def test_manual_watering_services(
         service_data={ATTR_ENTITY_ID: "switch.zone_one_manual_watering"},
         blocking=True,
     )
-    mock_pydrawise.run_zone.assert_called_once_with(15, 1)
+    mock_pydrawise.start_zone.assert_called_once_with(
+        zones[0], custom_run_duration=DEFAULT_WATERING_TIME
+    )
     state = hass.states.get("switch.zone_one_manual_watering")
     assert state is not None
     assert state.state == "on"
@@ -56,14 +64,18 @@ async def test_manual_watering_services(
         service_data={ATTR_ENTITY_ID: "switch.zone_one_manual_watering"},
         blocking=True,
     )
-    mock_pydrawise.run_zone.assert_called_once_with(0, 1)
+    mock_pydrawise.stop_zone.assert_called_once_with(zones[0])
     state = hass.states.get("switch.zone_one_manual_watering")
     assert state is not None
     assert state.state == "off"
 
 
+@pytest.mark.freeze_time("2023-10-01 00:00:00+00:00")
 async def test_auto_watering_services(
-    hass: HomeAssistant, mock_added_config_entry: MockConfigEntry, mock_pydrawise: Mock
+    hass: HomeAssistant,
+    mock_added_config_entry: MockConfigEntry,
+    mock_pydrawise: AsyncMock,
+    zones: list[Zone],
 ) -> None:
     """Test Automatic Watering services."""
     await hass.services.async_call(
@@ -72,7 +84,9 @@ async def test_auto_watering_services(
         service_data={ATTR_ENTITY_ID: "switch.zone_one_automatic_watering"},
         blocking=True,
     )
-    mock_pydrawise.suspend_zone.assert_called_once_with(365, 1)
+    mock_pydrawise.suspend_zone.assert_called_once_with(
+        zones[0], dt_util.now() + timedelta(days=365)
+    )
     state = hass.states.get("switch.zone_one_automatic_watering")
     assert state is not None
     assert state.state == "off"
@@ -84,7 +98,7 @@ async def test_auto_watering_services(
         service_data={ATTR_ENTITY_ID: "switch.zone_one_automatic_watering"},
         blocking=True,
     )
-    mock_pydrawise.suspend_zone.assert_called_once_with(0, 1)
+    mock_pydrawise.resume_zone.assert_called_once_with(zones[0])
     state = hass.states.get("switch.zone_one_automatic_watering")
     assert state is not None
     assert state.state == "on"