Add A. O. Smith integration (#104976)

CODEOWNERS

@@ -86,6 +86,8 @@ build.json @home-assistant/supervisor
 /tests/components/anova/ @Lash-L
 /homeassistant/components/anthemav/ @hyralex
 /tests/components/anthemav/ @hyralex
+/homeassistant/components/aosmith/ @bdr99
+/tests/components/aosmith/ @bdr99
 /homeassistant/components/apache_kafka/ @bachya
 /tests/components/apache_kafka/ @bachya
 /homeassistant/components/apcupsd/ @yuxincs

homeassistant/components/aosmith/__init__.py

@@ -0,0 +1,53 @@
+"""The A. O. Smith integration."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from py_aosmith import AOSmithAPIClient
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, Platform
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import aiohttp_client
+
+from .const import DOMAIN
+from .coordinator import AOSmithCoordinator
+
+PLATFORMS: list[Platform] = [Platform.WATER_HEATER]
+
+
+@dataclass
+class AOSmithData:
+    """Data for the A. O. Smith integration."""
+
+    coordinator: AOSmithCoordinator
+    client: AOSmithAPIClient
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Set up A. O. Smith from a config entry."""
+    email = entry.data[CONF_EMAIL]
+    password = entry.data[CONF_PASSWORD]
+
+    session = aiohttp_client.async_get_clientsession(hass)
+    client = AOSmithAPIClient(email, password, session)
+    coordinator = AOSmithCoordinator(hass, client)
+
+    # Fetch initial data so we have data when entities subscribe
+    await coordinator.async_config_entry_first_refresh()
+
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = AOSmithData(
+        coordinator=coordinator, client=client
+    )
+
+    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+
+    return True
+
+
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Unload a config entry."""
+    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+        hass.data[DOMAIN].pop(entry.entry_id)
+
+    return unload_ok

homeassistant/components/aosmith/config_flow.py

@@ -0,0 +1,61 @@
+"""Config flow for A. O. Smith integration."""
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from py_aosmith import AOSmithAPIClient, AOSmithInvalidCredentialsException
+import voluptuous as vol
+
+from homeassistant import config_entries
+from homeassistant.const import CONF_EMAIL, CONF_PASSWORD
+from homeassistant.data_entry_flow import FlowResult
+from homeassistant.helpers import aiohttp_client
+
+from .const import DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
+    """Handle a config flow for A. O. Smith."""
+
+    VERSION = 1
+
+    async def async_step_user(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Handle the initial step."""
+        errors: dict[str, str] = {}
+        if user_input is not None:
+            unique_id = user_input[CONF_EMAIL].lower()
+            await self.async_set_unique_id(unique_id)
+            self._abort_if_unique_id_configured()
+
+            session = aiohttp_client.async_get_clientsession(self.hass)
+            client = AOSmithAPIClient(
+                user_input[CONF_EMAIL], user_input[CONF_PASSWORD], session
+            )
+
+            try:
+                await client.get_devices()
+            except AOSmithInvalidCredentialsException:
+                errors["base"] = "invalid_auth"
+            except Exception:  # pylint: disable=broad-except
+                _LOGGER.exception("Unexpected exception")
+                errors["base"] = "unknown"
+            else:
+                return self.async_create_entry(
+                    title=user_input[CONF_EMAIL], data=user_input
+                )
+
+        return self.async_show_form(
+            step_id="user",
+            data_schema=vol.Schema(
+                {
+                    vol.Required(CONF_EMAIL): str,
+                    vol.Required(CONF_PASSWORD): str,
+                }
+            ),
+            errors=errors,
+        )

homeassistant/components/aosmith/const.py

@@ -0,0 +1,16 @@
+"""Constants for the A. O. Smith integration."""
+
+from datetime import timedelta
+
+DOMAIN = "aosmith"
+
+AOSMITH_MODE_ELECTRIC = "ELECTRIC"
+AOSMITH_MODE_HEAT_PUMP = "HEAT_PUMP"
+AOSMITH_MODE_HYBRID = "HYBRID"
+AOSMITH_MODE_VACATION = "VACATION"
+
+# Update interval to be used for normal background updates.
+REGULAR_INTERVAL = timedelta(seconds=30)
+
+# Update interval to be used while a mode or setpoint change is in progress.
+FAST_INTERVAL = timedelta(seconds=1)

homeassistant/components/aosmith/coordinator.py

@@ -0,0 +1,48 @@
+"""The data update coordinator for the A. O. Smith integration."""
+
+import logging
+from typing import Any
+
+from py_aosmith import (
+    AOSmithAPIClient,
+    AOSmithInvalidCredentialsException,
+    AOSmithUnknownException,
+)
+
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DOMAIN, FAST_INTERVAL, REGULAR_INTERVAL
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class AOSmithCoordinator(DataUpdateCoordinator[dict[str, dict[str, Any]]]):
+    """Custom data update coordinator for A. O. Smith integration."""
+
+    def __init__(self, hass: HomeAssistant, client: AOSmithAPIClient) -> None:
+        """Initialize the coordinator."""
+        super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=REGULAR_INTERVAL)
+        self.client = client
+
+    async def _async_update_data(self) -> dict[str, dict[str, Any]]:
+        """Fetch latest data from API."""
+        try:
+            devices = await self.client.get_devices()
+        except (AOSmithInvalidCredentialsException, AOSmithUnknownException) as err:
+            raise UpdateFailed(f"Error communicating with API: {err}") from err
+
+        mode_pending = any(
+            device.get("data", {}).get("modePending") for device in devices
+        )
+        setpoint_pending = any(
+            device.get("data", {}).get("temperatureSetpointPending")
+            for device in devices
+        )
+
+        if mode_pending or setpoint_pending:
+            self.update_interval = FAST_INTERVAL
+        else:
+            self.update_interval = REGULAR_INTERVAL
+
+        return {device.get("junctionId"): device for device in devices}

homeassistant/components/aosmith/entity.py

@@ -0,0 +1,51 @@
+"""The base entity for the A. O. Smith integration."""
+
+
+from py_aosmith import AOSmithAPIClient
+
+from homeassistant.helpers.device_registry import DeviceInfo
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+
+from .const import DOMAIN
+from .coordinator import AOSmithCoordinator
+
+
+class AOSmithEntity(CoordinatorEntity[AOSmithCoordinator]):
+    """Base entity for A. O. Smith."""
+
+    _attr_has_entity_name = True
+
+    def __init__(self, coordinator: AOSmithCoordinator, junction_id: str) -> None:
+        """Initialize the entity."""
+        super().__init__(coordinator)
+        self.junction_id = junction_id
+        self._attr_device_info = DeviceInfo(
+            manufacturer="A. O. Smith",
+            name=self.device.get("name"),
+            model=self.device.get("model"),
+            serial_number=self.device.get("serial"),
+            suggested_area=self.device.get("install", {}).get("location"),
+            identifiers={(DOMAIN, junction_id)},
+            sw_version=self.device.get("data", {}).get("firmwareVersion"),
+        )
+
+    @property
+    def device(self):
+        """Shortcut to get the device status from the coordinator data."""
+        return self.coordinator.data.get(self.junction_id)
+
+    @property
+    def device_data(self):
+        """Shortcut to get the device data within the device status."""
+        device = self.device
+        return None if device is None else device.get("data", {})
+
+    @property
+    def client(self) -> AOSmithAPIClient:
+        """Shortcut to get the API client."""
+        return self.coordinator.client
+
+    @property
+    def available(self) -> bool:
+        """Return True if entity is available."""
+        return super().available and self.device_data.get("isOnline") is True

homeassistant/components/aosmith/manifest.json

@@ -0,0 +1,10 @@
+{
+  "domain": "aosmith",
+  "name": "A. O. Smith",
+  "codeowners": ["@bdr99"],
+  "config_flow": true,
+  "documentation": "https://www.home-assistant.io/integrations/aosmith",
+  "iot_class": "cloud_polling",
+  "quality_scale": "platinum",
+  "requirements": ["py-aosmith==1.0.1"]
+}

homeassistant/components/aosmith/strings.json

@@ -0,0 +1,20 @@
+{
+  "config": {
+    "step": {
+      "user": {
+        "data": {
+          "email": "[%key:common::config_flow::data::email%]",
+          "password": "[%key:common::config_flow::data::password%]"
+        },
+        "description": "Please enter your A. O. Smith credentials."
+      }
+    },
+    "error": {
+      "invalid_auth": "[%key:common::config_flow::error::invalid_auth%]",
+      "unknown": "[%key:common::config_flow::error::unknown%]"
+    },
+    "abort": {
+      "already_configured": "[%key:common::config_flow::abort::already_configured_account%]"
+    }
+  }
+}

homeassistant/components/aosmith/water_heater.py

@@ -0,0 +1,149 @@
+"""The water heater platform for the A. O. Smith integration."""
+
+from typing import Any
+
+from homeassistant.components.water_heater import (
+    STATE_ECO,
+    STATE_ELECTRIC,
+    STATE_HEAT_PUMP,
+    STATE_OFF,
+    WaterHeaterEntity,
+    WaterHeaterEntityFeature,
+)
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import UnitOfTemperature
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from . import AOSmithData
+from .const import (
+    AOSMITH_MODE_ELECTRIC,
+    AOSMITH_MODE_HEAT_PUMP,
+    AOSMITH_MODE_HYBRID,
+    AOSMITH_MODE_VACATION,
+    DOMAIN,
+)
+from .coordinator import AOSmithCoordinator
+from .entity import AOSmithEntity
+
+MODE_HA_TO_AOSMITH = {
+    STATE_OFF: AOSMITH_MODE_VACATION,
+    STATE_ECO: AOSMITH_MODE_HYBRID,
+    STATE_ELECTRIC: AOSMITH_MODE_ELECTRIC,
+    STATE_HEAT_PUMP: AOSMITH_MODE_HEAT_PUMP,
+}
+MODE_AOSMITH_TO_HA = {
+    AOSMITH_MODE_ELECTRIC: STATE_ELECTRIC,
+    AOSMITH_MODE_HEAT_PUMP: STATE_HEAT_PUMP,
+    AOSMITH_MODE_HYBRID: STATE_ECO,
+    AOSMITH_MODE_VACATION: STATE_OFF,
+}
+
+# Operation mode to use when exiting away mode
+DEFAULT_OPERATION_MODE = AOSMITH_MODE_HYBRID
+
+DEFAULT_SUPPORT_FLAGS = (
+    WaterHeaterEntityFeature.TARGET_TEMPERATURE
+    | WaterHeaterEntityFeature.OPERATION_MODE
+)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+) -> None:
+    """Set up A. O. Smith water heater platform."""
+    data: AOSmithData = hass.data[DOMAIN][entry.entry_id]
+
+    entities = []
+
+    for junction_id in data.coordinator.data:
+        entities.append(AOSmithWaterHeaterEntity(data.coordinator, junction_id))
+
+    async_add_entities(entities)
+
+
+class AOSmithWaterHeaterEntity(AOSmithEntity, WaterHeaterEntity):
+    """The water heater entity for the A. O. Smith integration."""
+
+    _attr_name = None
+    _attr_temperature_unit = UnitOfTemperature.FAHRENHEIT
+    _attr_min_temp = 95
+
+    def __init__(self, coordinator: AOSmithCoordinator, junction_id: str) -> None:
+        """Initialize the entity."""
+        super().__init__(coordinator, junction_id)
+        self._attr_unique_id = junction_id
+
+    @property
+    def operation_list(self) -> list[str]:
+        """Return the list of supported operation modes."""
+        op_modes = []
+        for mode_dict in self.device_data.get("modes", []):
+            mode_name = mode_dict.get("mode")
+            ha_mode = MODE_AOSMITH_TO_HA.get(mode_name)
+
+            # Filtering out STATE_OFF since it is handled by away mode
+            if ha_mode is not None and ha_mode != STATE_OFF:
+                op_modes.append(ha_mode)
+
+        return op_modes
+
+    @property
+    def supported_features(self) -> WaterHeaterEntityFeature:
+        """Return the list of supported features."""
+        supports_vacation_mode = any(
+            mode_dict.get("mode") == AOSMITH_MODE_VACATION
+            for mode_dict in self.device_data.get("modes", [])
+        )
+
+        if supports_vacation_mode:
+            return DEFAULT_SUPPORT_FLAGS | WaterHeaterEntityFeature.AWAY_MODE
+
+        return DEFAULT_SUPPORT_FLAGS
+
+    @property
+    def target_temperature(self) -> float | None:
+        """Return the temperature we try to reach."""
+        return self.device_data.get("temperatureSetpoint")
+
+    @property
+    def max_temp(self) -> float:
+        """Return the maximum temperature."""
+        return self.device_data.get("temperatureSetpointMaximum")
+
+    @property
+    def current_operation(self) -> str:
+        """Return the current operation mode."""
+        return MODE_AOSMITH_TO_HA.get(self.device_data.get("mode"), STATE_OFF)
+
+    @property
+    def is_away_mode_on(self):
+        """Return True if away mode is on."""
+        return self.device_data.get("mode") == AOSMITH_MODE_VACATION
+
+    async def async_set_operation_mode(self, operation_mode: str) -> None:
+        """Set new target operation mode."""
+        aosmith_mode = MODE_HA_TO_AOSMITH.get(operation_mode)
+        if aosmith_mode is not None:
+            await self.client.update_mode(self.junction_id, aosmith_mode)
+
+        await self.coordinator.async_request_refresh()
+
+    async def async_set_temperature(self, **kwargs: Any) -> None:
+        """Set new target temperature."""
+        temperature = kwargs.get("temperature")
+        await self.client.update_setpoint(self.junction_id, temperature)
+
+        await self.coordinator.async_request_refresh()
+
+    async def async_turn_away_mode_on(self) -> None:
+        """Turn away mode on."""
+        await self.client.update_mode(self.junction_id, AOSMITH_MODE_VACATION)
+
+        await self.coordinator.async_request_refresh()
+
+    async def async_turn_away_mode_off(self) -> None:
+        """Turn away mode off."""
+        await self.client.update_mode(self.junction_id, DEFAULT_OPERATION_MODE)
+
+        await self.coordinator.async_request_refresh()

homeassistant/generated/config_flows.py

@@ -46,6 +46,7 @@ FLOWS = {
         "androidtv_remote",
         "anova",
         "anthemav",
+        "aosmith",
         "apcupsd",
         "apple_tv",
         "aranet",

homeassistant/generated/integrations.json

@@ -286,6 +286,12 @@
       "integration_type": "virtual",
       "supported_by": "energyzero"
     },
+    "aosmith": {
+      "name": "A. O. Smith",
+      "integration_type": "hub",
+      "config_flow": true,
+      "iot_class": "cloud_polling"
+    },
     "apache_kafka": {
       "name": "Apache Kafka",
       "integration_type": "hub",

requirements_all.txt

@@ -1534,6 +1534,9 @@ pushover_complete==1.1.1
 # homeassistant.components.pvoutput
 pvo==2.1.1
 
+# homeassistant.components.aosmith
+py-aosmith==1.0.1
+
 # homeassistant.components.canary
 py-canary==0.5.3
 

requirements_test_all.txt

@@ -1177,6 +1177,9 @@ pushover_complete==1.1.1
 # homeassistant.components.pvoutput
 pvo==2.1.1
 
+# homeassistant.components.aosmith
+py-aosmith==1.0.1
+
 # homeassistant.components.canary
 py-canary==0.5.3
 

tests/components/aosmith/__init__.py

@@ -0,0 +1 @@
+"""Tests for the A. O. Smith integration."""

tests/components/aosmith/conftest.py

@@ -0,0 +1,74 @@
+"""Common fixtures for the A. O. Smith tests."""
+from collections.abc import Generator
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from py_aosmith import AOSmithAPIClient
+import pytest
+
+from homeassistant.components.aosmith.const import DOMAIN
+from homeassistant.const import CONF_EMAIL, CONF_PASSWORD
+from homeassistant.core import HomeAssistant
+from homeassistant.util.unit_system import US_CUSTOMARY_SYSTEM
+
+from tests.common import MockConfigEntry, load_json_array_fixture
+
+FIXTURE_USER_INPUT = {
+    CONF_EMAIL: "testemail@example.com",
+    CONF_PASSWORD: "test-password",
+}
+
+
+@pytest.fixture
+def mock_config_entry() -> MockConfigEntry:
+    """Return the default mocked config entry."""
+    return MockConfigEntry(
+        domain=DOMAIN,
+        data=FIXTURE_USER_INPUT,
+        unique_id="unique_id",
+    )
+
+
+@pytest.fixture
+def mock_setup_entry() -> Generator[AsyncMock, None, None]:
+    """Override async_setup_entry."""
+    with patch(
+        "homeassistant.components.aosmith.async_setup_entry", return_value=True
+    ) as mock_setup_entry:
+        yield mock_setup_entry
+
+
+@pytest.fixture
+def get_devices_fixture() -> str:
+    """Return the name of the fixture to use for get_devices."""
+    return "get_devices"
+
+
+@pytest.fixture
+async def mock_client(get_devices_fixture: str) -> Generator[MagicMock, None, None]:
+    """Return a mocked client."""
+    get_devices_fixture = load_json_array_fixture(f"{get_devices_fixture}.json", DOMAIN)
+
+    client_mock = MagicMock(AOSmithAPIClient)
+    client_mock.get_devices = AsyncMock(return_value=get_devices_fixture)
+
+    return client_mock
+
+
+@pytest.fixture
+async def init_integration(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_client: MagicMock,
+) -> MockConfigEntry:
+    """Set up the integration for testing."""
+    hass.config.units = US_CUSTOMARY_SYSTEM
+
+    with patch(
+        "homeassistant.components.aosmith.AOSmithAPIClient", return_value=mock_client
+    ):
+        mock_config_entry.add_to_hass(hass)
+
+        await hass.config_entries.async_setup(mock_config_entry.entry_id)
+        await hass.async_block_till_done()
+
+        return mock_config_entry

tests/components/aosmith/fixtures/get_devices.json

@@ -0,0 +1,46 @@
+[
+  {
+    "brand": "aosmith",
+    "model": "HPTS-50 200 202172000",
+    "deviceType": "NEXT_GEN_HEAT_PUMP",
+    "dsn": "dsn",
+    "junctionId": "junctionId",
+    "name": "My water heater",
+    "serial": "serial",
+    "install": {
+      "location": "Basement"
+    },
+    "data": {
+      "__typename": "NextGenHeatPump",
+      "temperatureSetpoint": 130,
+      "temperatureSetpointPending": false,
+      "temperatureSetpointPrevious": 130,
+      "temperatureSetpointMaximum": 130,
+      "modes": [
+        {
+          "mode": "HYBRID",
+          "controls": null
+        },
+        {
+          "mode": "HEAT_PUMP",
+          "controls": null
+        },
+        {
+          "mode": "ELECTRIC",
+          "controls": "SELECT_DAYS"
+        },
+        {
+          "mode": "VACATION",
+          "controls": "SELECT_DAYS"
+        }
+      ],
+      "isOnline": true,
+      "firmwareVersion": "2.14",
+      "hotWaterStatus": "LOW",
+      "mode": "HEAT_PUMP",
+      "modePending": false,
+      "vacationModeRemainingDays": 0,
+      "electricModeRemainingDays": 0
+    }
+  }
+]

tests/components/aosmith/fixtures/get_devices_mode_pending.json

@@ -0,0 +1,46 @@
+[
+  {
+    "brand": "aosmith",
+    "model": "HPTS-50 200 202172000",
+    "deviceType": "NEXT_GEN_HEAT_PUMP",
+    "dsn": "dsn",
+    "junctionId": "junctionId",
+    "name": "My water heater",
+    "serial": "serial",
+    "install": {
+      "location": "Basement"
+    },
+    "data": {
+      "__typename": "NextGenHeatPump",
+      "temperatureSetpoint": 130,
+      "temperatureSetpointPending": false,
+      "temperatureSetpointPrevious": 130,
+      "temperatureSetpointMaximum": 130,
+      "modes": [
+        {
+          "mode": "HYBRID",
+          "controls": null
+        },
+        {
+          "mode": "HEAT_PUMP",
+          "controls": null
+        },
+        {
+          "mode": "ELECTRIC",
+          "controls": "SELECT_DAYS"
+        },
+        {
+          "mode": "VACATION",
+          "controls": "SELECT_DAYS"
+        }
+      ],
+      "isOnline": true,
+      "firmwareVersion": "2.14",
+      "hotWaterStatus": "LOW",
+      "mode": "HEAT_PUMP",
+      "modePending": true,
+      "vacationModeRemainingDays": 0,
+      "electricModeRemainingDays": 0
+    }
+  }
+]

tests/components/aosmith/fixtures/get_devices_no_vacation_mode.json

@@ -0,0 +1,42 @@
+[
+  {
+    "brand": "aosmith",
+    "model": "HPTS-50 200 202172000",
+    "deviceType": "NEXT_GEN_HEAT_PUMP",
+    "dsn": "dsn",
+    "junctionId": "junctionId",
+    "name": "My water heater",
+    "serial": "serial",
+    "install": {
+      "location": "Basement"
+    },
+    "data": {
+      "__typename": "NextGenHeatPump",
+      "temperatureSetpoint": 130,
+      "temperatureSetpointPending": false,
+      "temperatureSetpointPrevious": 130,
+      "temperatureSetpointMaximum": 130,
+      "modes": [
+        {
+          "mode": "HYBRID",
+          "controls": null
+        },
+        {
+          "mode": "HEAT_PUMP",
+          "controls": null
+        },
+        {
+          "mode": "ELECTRIC",
+          "controls": "SELECT_DAYS"
+        }
+      ],
+      "isOnline": true,
+      "firmwareVersion": "2.14",
+      "hotWaterStatus": "LOW",
+      "mode": "HEAT_PUMP",
+      "modePending": false,
+      "vacationModeRemainingDays": 0,
+      "electricModeRemainingDays": 0
+    }
+  }
+]

tests/components/aosmith/fixtures/get_devices_setpoint_pending.json

@@ -0,0 +1,46 @@
+[
+  {
+    "brand": "aosmith",
+    "model": "HPTS-50 200 202172000",
+    "deviceType": "NEXT_GEN_HEAT_PUMP",
+    "dsn": "dsn",
+    "junctionId": "junctionId",
+    "name": "My water heater",
+    "serial": "serial",
+    "install": {
+      "location": "Basement"
+    },
+    "data": {
+      "__typename": "NextGenHeatPump",
+      "temperatureSetpoint": 130,
+      "temperatureSetpointPending": true,
+      "temperatureSetpointPrevious": 130,
+      "temperatureSetpointMaximum": 130,
+      "modes": [
+        {
+          "mode": "HYBRID",
+          "controls": null
+        },
+        {
+          "mode": "HEAT_PUMP",
+          "controls": null
+        },
+        {
+          "mode": "ELECTRIC",
+          "controls": "SELECT_DAYS"
+        },
+        {
+          "mode": "VACATION",
+          "controls": "SELECT_DAYS"
+        }
+      ],
+      "isOnline": true,
+      "firmwareVersion": "2.14",
+      "hotWaterStatus": "LOW",
+      "mode": "HEAT_PUMP",
+      "modePending": false,
+      "vacationModeRemainingDays": 0,
+      "electricModeRemainingDays": 0
+    }
+  }
+]

tests/components/aosmith/snapshots/test_device.ambr

@@ -0,0 +1,29 @@
+# serializer version: 1
+# name: test_device
+  DeviceRegistryEntrySnapshot({
+    'area_id': 'basement',
+    'config_entries': <ANY>,
+    'configuration_url': None,
+    'connections': set({
+    }),
+    'disabled_by': None,
+    'entry_type': None,
+    'hw_version': None,
+    'id': <ANY>,
+    'identifiers': set({
+      tuple(
+        'aosmith',
+        'junctionId',
+      ),
+    }),
+    'is_new': False,
+    'manufacturer': 'A. O. Smith',
+    'model': 'HPTS-50 200 202172000',
+    'name': 'My water heater',
+    'name_by_user': None,
+    'serial_number': 'serial',
+    'suggested_area': 'Basement',
+    'sw_version': '2.14',
+    'via_device_id': None,
+  })
+# ---

tests/components/aosmith/snapshots/test_water_heater.ambr

@@ -0,0 +1,27 @@
+# serializer version: 1
+# name: test_state
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'away_mode': 'off',
+      'current_temperature': None,
+      'friendly_name': 'My water heater',
+      'max_temp': 130,
+      'min_temp': 95,
+      'operation_list': list([
+        'eco',
+        'heat_pump',
+        'electric',
+      ]),
+      'operation_mode': 'heat_pump',
+      'supported_features': <WaterHeaterEntityFeature: 7>,
+      'target_temp_high': None,
+      'target_temp_low': None,
+      'temperature': 130,
+    }),
+    'context': <ANY>,
+    'entity_id': 'water_heater.my_water_heater',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'heat_pump',
+  })
+# ---

tests/components/aosmith/test_config_flow.py

@@ -0,0 +1,84 @@
+"""Test the A. O. Smith config flow."""
+from unittest.mock import AsyncMock, patch
+
+from py_aosmith import AOSmithInvalidCredentialsException
+import pytest
+
+from homeassistant import config_entries
+from homeassistant.components.aosmith.const import DOMAIN
+from homeassistant.const import CONF_EMAIL
+from homeassistant.core import HomeAssistant
+from homeassistant.data_entry_flow import FlowResultType
+
+from tests.components.aosmith.conftest import FIXTURE_USER_INPUT
+
+
+async def test_form(hass: HomeAssistant, mock_setup_entry: AsyncMock) -> None:
+    """Test we get the form."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    assert result["type"] == FlowResultType.FORM
+    assert result["errors"] == {}
+
+    with patch(
+        "homeassistant.components.aosmith.config_flow.AOSmithAPIClient.get_devices",
+        return_value=[],
+    ):
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            FIXTURE_USER_INPUT,
+        )
+        await hass.async_block_till_done()
+
+    assert result2["type"] == FlowResultType.CREATE_ENTRY
+    assert result2["title"] == FIXTURE_USER_INPUT[CONF_EMAIL]
+    assert result2["data"] == FIXTURE_USER_INPUT
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
+@pytest.mark.parametrize(
+    ("exception", "expected_error_key"),
+    [
+        (AOSmithInvalidCredentialsException("Invalid credentials"), "invalid_auth"),
+        (Exception, "unknown"),
+    ],
+)
+async def test_form_exception(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    exception: Exception,
+    expected_error_key: str,
+) -> None:
+    """Test handling an exception and then recovering on the second attempt."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    assert result["type"] == FlowResultType.FORM
+    assert result["errors"] == {}
+
+    with patch(
+        "homeassistant.components.aosmith.config_flow.AOSmithAPIClient.get_devices",
+        side_effect=exception,
+    ):
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            FIXTURE_USER_INPUT,
+        )
+        assert result2["type"] == FlowResultType.FORM
+        assert result2["errors"] == {"base": expected_error_key}
+
+    with patch(
+        "homeassistant.components.aosmith.config_flow.AOSmithAPIClient.get_devices",
+        return_value=[],
+    ):
+        result3 = await hass.config_entries.flow.async_configure(
+            result2["flow_id"],
+            FIXTURE_USER_INPUT,
+        )
+        await hass.async_block_till_done()
+
+    assert result3["type"] == FlowResultType.CREATE_ENTRY
+    assert result3["title"] == FIXTURE_USER_INPUT[CONF_EMAIL]
+    assert result3["data"] == FIXTURE_USER_INPUT
+    assert len(mock_setup_entry.mock_calls) == 1

tests/components/aosmith/test_device.py

@@ -0,0 +1,23 @@
+"""Tests for the device created by the A. O. Smith integration."""
+
+from syrupy.assertion import SnapshotAssertion
+
+from homeassistant.components.aosmith.const import DOMAIN
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr
+
+from tests.common import MockConfigEntry
+
+
+async def test_device(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    init_integration: MockConfigEntry,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test creation of the device."""
+    reg_device = device_registry.async_get_device(
+        identifiers={(DOMAIN, "junctionId")},
+    )
+
+    assert reg_device == snapshot

tests/components/aosmith/test_init.py

@@ -0,0 +1,71 @@
+"""Tests for the initialization of the A. O. Smith integration."""
+
+from datetime import timedelta
+from unittest.mock import MagicMock, patch
+
+from freezegun.api import FrozenDateTimeFactory
+from py_aosmith import AOSmithUnknownException
+import pytest
+
+from homeassistant.components.aosmith.const import (
+    DOMAIN,
+    FAST_INTERVAL,
+    REGULAR_INTERVAL,
+)
+from homeassistant.config_entries import ConfigEntryState
+from homeassistant.core import HomeAssistant
+
+from tests.common import MockConfigEntry, async_fire_time_changed
+
+
+async def test_config_entry_setup(init_integration: MockConfigEntry) -> None:
+    """Test setup of the config entry."""
+    mock_config_entry = init_integration
+
+    assert mock_config_entry.state is ConfigEntryState.LOADED
+
+
+async def test_config_entry_not_ready(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
+    """Test the config entry not ready."""
+    mock_config_entry.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.aosmith.config_flow.AOSmithAPIClient.get_devices",
+        side_effect=AOSmithUnknownException("Unknown error"),
+    ):
+        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
+
+
+@pytest.mark.parametrize(
+    ("get_devices_fixture", "time_to_wait", "expected_call_count"),
+    [
+        ("get_devices", REGULAR_INTERVAL, 1),
+        ("get_devices", FAST_INTERVAL, 0),
+        ("get_devices_mode_pending", FAST_INTERVAL, 1),
+        ("get_devices_setpoint_pending", FAST_INTERVAL, 1),
+    ],
+)
+async def test_update(
+    freezer: FrozenDateTimeFactory,
+    hass: HomeAssistant,
+    mock_client: MagicMock,
+    init_integration: MockConfigEntry,
+    time_to_wait: timedelta,
+    expected_call_count: int,
+) -> None:
+    """Test data update with differing intervals depending on device status."""
+    entries = hass.config_entries.async_entries(DOMAIN)
+    assert len(entries) == 1
+    assert entries[0].state is ConfigEntryState.LOADED
+    assert mock_client.get_devices.call_count == 1
+
+    freezer.tick(time_to_wait)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    assert mock_client.get_devices.call_count == 1 + expected_call_count

tests/components/aosmith/test_water_heater.py

@@ -0,0 +1,147 @@
+"""Tests for the water heater platform of the A. O. Smith integration."""
+
+from unittest.mock import MagicMock
+
+import pytest
+from syrupy.assertion import SnapshotAssertion
+
+from homeassistant.components.aosmith.const import (
+    AOSMITH_MODE_ELECTRIC,
+    AOSMITH_MODE_HEAT_PUMP,
+    AOSMITH_MODE_HYBRID,
+    AOSMITH_MODE_VACATION,
+)
+from homeassistant.components.water_heater import (
+    ATTR_AWAY_MODE,
+    ATTR_OPERATION_MODE,
+    ATTR_TEMPERATURE,
+    DOMAIN as WATER_HEATER_DOMAIN,
+    SERVICE_SET_AWAY_MODE,
+    SERVICE_SET_OPERATION_MODE,
+    SERVICE_SET_TEMPERATURE,
+    STATE_ECO,
+    STATE_ELECTRIC,
+    STATE_HEAT_PUMP,
+    WaterHeaterEntityFeature,
+)
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    ATTR_FRIENDLY_NAME,
+    ATTR_SUPPORTED_FEATURES,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
+
+from tests.common import MockConfigEntry
+
+
+async def test_setup(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    init_integration: MockConfigEntry,
+) -> None:
+    """Test the setup of the water heater entity."""
+    entry = entity_registry.async_get("water_heater.my_water_heater")
+    assert entry
+    assert entry.unique_id == "junctionId"
+
+    state = hass.states.get("water_heater.my_water_heater")
+    assert state
+    assert state.attributes.get(ATTR_FRIENDLY_NAME) == "My water heater"
+
+
+async def test_state(
+    hass: HomeAssistant, init_integration: MockConfigEntry, snapshot: SnapshotAssertion
+) -> None:
+    """Test the state of the water heater entity."""
+    state = hass.states.get("water_heater.my_water_heater")
+    assert state == snapshot
+
+
+@pytest.mark.parametrize(
+    ("get_devices_fixture"),
+    ["get_devices_no_vacation_mode"],
+)
+async def test_state_away_mode_unsupported(
+    hass: HomeAssistant, init_integration: MockConfigEntry
+) -> None:
+    """Test that away mode is not supported if the water heater does not support vacation mode."""
+    state = hass.states.get("water_heater.my_water_heater")
+    assert (
+        state.attributes.get(ATTR_SUPPORTED_FEATURES)
+        == WaterHeaterEntityFeature.TARGET_TEMPERATURE
+        | WaterHeaterEntityFeature.OPERATION_MODE
+    )
+
+
+@pytest.mark.parametrize(
+    ("hass_mode", "aosmith_mode"),
+    [
+        (STATE_HEAT_PUMP, AOSMITH_MODE_HEAT_PUMP),
+        (STATE_ECO, AOSMITH_MODE_HYBRID),
+        (STATE_ELECTRIC, AOSMITH_MODE_ELECTRIC),
+    ],
+)
+async def test_set_operation_mode(
+    hass: HomeAssistant,
+    mock_client: MagicMock,
+    init_integration: MockConfigEntry,
+    hass_mode: str,
+    aosmith_mode: str,
+) -> None:
+    """Test setting the operation mode."""
+    await hass.services.async_call(
+        WATER_HEATER_DOMAIN,
+        SERVICE_SET_OPERATION_MODE,
+        {
+            ATTR_ENTITY_ID: "water_heater.my_water_heater",
+            ATTR_OPERATION_MODE: hass_mode,
+        },
+    )
+    await hass.async_block_till_done()
+
+    mock_client.update_mode.assert_called_once_with("junctionId", aosmith_mode)
+
+
+async def test_set_temperature(
+    hass: HomeAssistant,
+    mock_client: MagicMock,
+    init_integration: MockConfigEntry,
+) -> None:
+    """Test setting the target temperature."""
+    await hass.services.async_call(
+        WATER_HEATER_DOMAIN,
+        SERVICE_SET_TEMPERATURE,
+        {ATTR_ENTITY_ID: "water_heater.my_water_heater", ATTR_TEMPERATURE: 120},
+    )
+    await hass.async_block_till_done()
+
+    mock_client.update_setpoint.assert_called_once_with("junctionId", 120)
+
+
+@pytest.mark.parametrize(
+    ("hass_away_mode", "aosmith_mode"),
+    [
+        (True, AOSMITH_MODE_VACATION),
+        (False, AOSMITH_MODE_HYBRID),
+    ],
+)
+async def test_away_mode(
+    hass: HomeAssistant,
+    mock_client: MagicMock,
+    init_integration: MockConfigEntry,
+    hass_away_mode: bool,
+    aosmith_mode: str,
+) -> None:
+    """Test turning away mode on/off."""
+    await hass.services.async_call(
+        WATER_HEATER_DOMAIN,
+        SERVICE_SET_AWAY_MODE,
+        {
+            ATTR_ENTITY_ID: "water_heater.my_water_heater",
+            ATTR_AWAY_MODE: hass_away_mode,
+        },
+    )
+    await hass.async_block_till_done()
+
+    mock_client.update_mode.assert_called_once_with("junctionId", aosmith_mode)