Add Mammotion integration

.strict-typing

@@ -370,6 +370,7 @@ homeassistant.components.luftdaten.*
 homeassistant.components.lunatone.*
 homeassistant.components.lutron.*
 homeassistant.components.madvr.*
+homeassistant.components.mammotion.*
 homeassistant.components.manual.*
 homeassistant.components.marantz_infrared.*
 homeassistant.components.mastodon.*

homeassistant/components/mammotion/__init__.py

@@ -0,0 +1,217 @@
+"""The Mammotion integration."""
+
+import contextlib
+from typing import Any
+
+from aiohttp import ClientConnectorError
+from pymammotion.aliyun.model.dev_by_account_response import Device
+from pymammotion.client import MammotionClient
+from pymammotion.homeassistant import HomeAssistantMowerApi
+from pymammotion.transport.base import LoginFailedError, TransportType
+from Tea.exceptions import UnretryableException
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import (
+    CONF_PASSWORD,
+    EVENT_HOMEASSISTANT_FINAL_WRITE,
+    EVENT_HOMEASSISTANT_STOP,
+    Platform,
+)
+from homeassistant.core import Event, HomeAssistant
+from homeassistant.exceptions import (
+    ConfigEntryAuthFailed,
+    ConfigEntryError,
+    ConfigEntryNotReady,
+)
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.device_registry import DeviceEntry
+
+from .config import MammotionConfigStore
+from .const import (
+    CONF_ACCOUNTNAME,
+    CONF_AEP_DATA,
+    CONF_MAMMOTION_DEVICE_RECORDS,
+    CONF_MAMMOTION_MQTT,
+    DEVICE_SUPPORT,
+    DOMAIN,
+    EXPIRED_CREDENTIAL_EXCEPTIONS,
+    LOGGER,
+)
+from .coordinator import MammotionMowerUpdateCoordinator
+from .models import MammotionDevices, MammotionMowerData
+
+PLATFORMS: list[Platform] = [Platform.LAWN_MOWER]
+
+type MammotionConfigEntry = ConfigEntry[MammotionDevices]
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: MammotionConfigEntry) -> bool:
+    """Set up Mammotion Luba from a config entry."""
+
+    api = HomeAssistantMowerApi(
+        ha_version="1.0.0", session=async_get_clientsession(hass)
+    )
+    mammotion = api.mammotion
+    account = entry.data.get(CONF_ACCOUNTNAME)
+    password = entry.data.get(CONF_PASSWORD)
+
+    mammotion_mowers: list[MammotionMowerData] = []
+    mammotion_devices: MammotionDevices = MammotionDevices([])
+    store = MammotionConfigStore(hass, entry.entry_id)
+    await store.async_load_mower_data()
+
+    if account and password:
+        session = async_get_clientsession(hass)
+        cached = _load_cached_credentials(entry)
+
+        async def credentials_updated() -> None:
+            """Persist credentials the library rotated behind our back."""
+            store_cloud_credentials(hass, entry, mammotion)
+
+        async def unrecoverable_auth_error(
+            account_id: str, transport_type: TransportType, err: Exception
+        ) -> None:
+            """Surface a dead cloud session as a single reauth prompt."""
+            LOGGER.error(
+                "Mammotion cloud authentication failed for %s on %s: %s",
+                account_id,
+                transport_type.value,
+                err,
+            )
+            entry.async_start_reauth(hass)
+
+        # Wired before login so a rotation during it is not lost.
+        mammotion.on_credentials_updated = credentials_updated
+        mammotion.on_unrecoverable_auth_error = unrecoverable_auth_error
+
+        try:
+            if cached:
+                try:
+                    await mammotion.restore_credentials(
+                        account, password, cached, session
+                    )
+                except EXPIRED_CREDENTIAL_EXCEPTIONS:
+                    await mammotion.login_and_initiate_cloud(account, password, session)
+            else:
+                await mammotion.login_and_initiate_cloud(account, password, session)
+        except ClientConnectorError as err:
+            raise ConfigEntryNotReady(err) from err
+        except LoginFailedError as err:
+            raise ConfigEntryAuthFailed(err) from err
+        except UnretryableException as err:
+            raise ConfigEntryError(err) from err
+
+        store_cloud_credentials(hass, entry, mammotion)
+
+        device_list: list[Device] = [
+            device
+            for device in (
+                *mammotion.aliyun_device_list,
+                *mammotion.mammotion_device_list,
+            )
+            if device.device_name.startswith(DEVICE_SUPPORT)
+        ]
+
+        for device in device_list:
+            update_coordinator = MammotionMowerUpdateCoordinator(
+                hass, entry, device, api, store
+            )
+
+            update_coordinator.restore_data()
+            await update_coordinator.async_config_entry_first_refresh()
+
+            mammotion_mowers.append(
+                MammotionMowerData(
+                    name=device.device_name,
+                    api=api,
+                    coordinator=update_coordinator,
+                    device=device,
+                )
+            )
+
+    mammotion_devices.mowers = mammotion_mowers
+    entry.runtime_data = mammotion_devices
+
+    async def shutdown_mammotion(_: Event | None = None) -> None:
+        await mammotion.stop()
+
+    entry.async_on_unload(
+        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, shutdown_mammotion)
+    )
+
+    entry.async_on_unload(
+        hass.bus.async_listen_once(
+            EVENT_HOMEASSISTANT_FINAL_WRITE, store.async_flush_mower_data
+        )
+    )
+
+    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+
+    return True
+
+
+def store_cloud_credentials(
+    hass: HomeAssistant,
+    config_entry: MammotionConfigEntry,
+    mammotion: MammotionClient,
+) -> None:
+    """Store cloud credentials in config entry."""
+    cache = mammotion.to_cache()
+    if not cache:
+        return
+    hass.config_entries.async_update_entry(
+        config_entry, data={**config_entry.data, **cache}
+    )
+
+
+def _load_cached_credentials(entry: MammotionConfigEntry) -> dict[str, Any]:
+    """Return the config entry's cached credential data, keyed as the library expects."""
+    has_aliyun = bool(entry.data.get(CONF_AEP_DATA))
+    has_mammotion = bool(entry.data.get(CONF_MAMMOTION_MQTT)) and bool(
+        entry.data.get(CONF_MAMMOTION_DEVICE_RECORDS)
+    )
+    return dict(entry.data) if (has_aliyun or has_mammotion) else {}
+
+
+async def async_unload_entry(hass: HomeAssistant, entry: MammotionConfigEntry) -> bool:
+    """Unload a config entry."""
+
+    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+        for mower in entry.runtime_data.mowers:
+            mower.coordinator.store_cloud_credentials()
+            with contextlib.suppress(TimeoutError):
+                await mower.api.mammotion.remove_device(mower.name)
+
+        if mowers := entry.runtime_data.mowers:
+            await mowers[0].coordinator.store.async_flush_mower_data()
+            await mowers[0].api.mammotion.stop()
+    return unload_ok
+
+
+async def async_remove_entry(hass: HomeAssistant, entry: MammotionConfigEntry) -> None:
+    """Remove a config entry."""
+    store = MammotionConfigStore(hass, entry.entry_id)
+    await store.async_remove()
+
+
+async def async_remove_config_entry_device(
+    hass: HomeAssistant, config_entry: MammotionConfigEntry, device_entry: DeviceEntry
+) -> bool:
+    """Remove a config entry from a device."""
+    mower_names = (
+        next(
+            identifier[1]
+            for identifier in device_entry.identifiers
+            if identifier[0] == DOMAIN
+        ),
+    )
+    mower = next(
+        (
+            mower
+            for mower in config_entry.runtime_data.mowers
+            if mower.name in mower_names
+        ),
+        None,
+    )
+
+    return not bool(mower)

homeassistant/components/mammotion/config.py

@@ -0,0 +1,45 @@
+"""Config storage for Mammotion integration."""
+
+from typing import Any
+
+from homeassistant.core import Event, HomeAssistant, callback
+from homeassistant.helpers.storage import Store
+
+from .const import DOMAIN
+
+STORAGE_VERSION = 1
+STORAGE_MINOR_VERSION = 0
+
+
+class MammotionConfigStore(Store[dict[str, Any]]):
+    """Store the mower state of a single config entry."""
+
+    def __init__(self, hass: HomeAssistant, entry_id: str) -> None:
+        """Initialize the configuration store."""
+        super().__init__(
+            hass,
+            version=STORAGE_VERSION,
+            minor_version=STORAGE_MINOR_VERSION,
+            key=f"{DOMAIN}.{entry_id}",
+        )
+        self.mower_data: dict[str, Any] = {}
+        self._dirty = False
+
+    async def async_load_mower_data(self) -> None:
+        """Load the persisted mower data into memory."""
+        self.mower_data = await self.async_load() or {}
+
+    @callback
+    def async_update_mower_data(self, device_name: str, data: dict[str, Any]) -> None:
+        """Update a mower in memory; the data is only written on shutdown."""
+        if self.mower_data.get(device_name) == data:
+            return
+        self.mower_data[device_name] = data
+        self._dirty = True
+
+    async def async_flush_mower_data(self, _event: Event | None = None) -> None:
+        """Write the in-memory mower data to disk if it changed."""
+        if not self._dirty:
+            return
+        self._dirty = False
+        await self.async_save(dict(self.mower_data))