add RTK devices and add relocate charging station

Author: mikey0000

custom_components/mammotion/__init__.py

@@ -13,7 +13,10 @@
 from pymammotion import CloudIOTGateway
 from pymammotion.aliyun.model.aep_response import AepResponse
 from pymammotion.aliyun.model.connect_response import ConnectResponse
-from pymammotion.aliyun.model.dev_by_account_response import ListingDevByAccountResponse
+from pymammotion.aliyun.model.dev_by_account_response import (
+    Device,
+    ListingDevByAccountResponse,
+)
 from pymammotion.aliyun.model.login_by_oauth_response import LoginByOAuthResponse
 from pymammotion.aliyun.model.regions_response import RegionResponse
 from pymammotion.aliyun.model.session_by_authcode_response import (
@@ -51,8 +54,9 @@
     MammotionMaintenanceUpdateCoordinator,
     MammotionMapUpdateCoordinator,
     MammotionReportUpdateCoordinator,
+    MammotionRTKCoordinator,
 )
-from .models import MammotionMowerData
+from .models import MammotionDevices, MammotionMowerData, MammotionRTKData
 
 PLATFORMS: list[Platform] = [
     Platform.BINARY_SENSOR,
@@ -67,7 +71,7 @@
     Platform.UPDATE,
 ]
 
-type MammotionConfigEntry = ConfigEntry[list[MammotionMowerData]]
+type MammotionConfigEntry = ConfigEntry[MammotionDevices]
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: MammotionConfigEntry) -> bool:
@@ -90,7 +94,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: MammotionConfigEntry) ->
 
     use_wifi = entry.data.get(CONF_USE_WIFI, True)
 
-    mammotion_devices: list[MammotionMowerData] = []
+    mammotion_mowers: list[MammotionMowerData] = []
+    mammotion_devices: MammotionDevices = MammotionDevices([], [])
+    mammotion_rtk: list[MammotionRTKData] = []
+    mammotion_rtk_devices: list[Device] = []
 
     if account and password:
         credentials = Credentials()
@@ -121,6 +128,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: MammotionConfigEntry) ->
                 device
             ) in mqtt_client.cloud_client.devices_by_account_response.data.data:
                 if not device.deviceName.startswith(DEVICE_SUPPORT):
+                    if device.categoryKey == "Tracker":
+                        mammotion_rtk_devices.append(device)
                     continue
 
                 mammotion_device = mammotion.get_or_create_device_by_name(
@@ -178,7 +187,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: MammotionConfigEntry) ->
                     # not entirely sure this is a good idea
                     mammotion_device.remove_cloud()
 
-                mammotion_devices.append(
+                mammotion_mowers.append(
                     MammotionMowerData(
                         name=device.deviceName,
                         device=device,
@@ -196,13 +205,26 @@ async def async_setup_entry(hass: HomeAssistant, entry: MammotionConfigEntry) ->
                 except:
                     """Do nothing for now."""
 
+            for rtk in mammotion_rtk_devices:
+                rtk_coordinator = MammotionRTKCoordinator(hass, entry, rtk, mqtt_client)
+                await rtk_coordinator.async_config_entry_first_refresh()
+                mammotion_rtk.append(
+                    MammotionRTKData(
+                        name=rtk.deviceName,
+                        api=mammotion,
+                        device=rtk,
+                        coordinator=rtk_coordinator,
+                    )
+                )
+
     # if not any(mammotion.get_device_by_name(mammotion_device.device.deviceName).preference == ConnectionPreference.WIFI for mammotion_device in mammotion_devices):
     #     for mammotion_device in mammotion_devices:
     #         mower = mammotion.get_device_by_name(mammotion_device.device.deviceName)
     #         await mower.cloud().stop()
     #         mower.cloud().mqtt.disconnect() if mower.cloud().mqtt.is_connected() else None
     #         mower.remove_cloud()
-
+    mammotion_devices.RTK = mammotion_rtk
+    mammotion_devices.mowers = mammotion_mowers
     entry.runtime_data = mammotion_devices
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -343,7 +365,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: MammotionConfigEntry) -
     """Unload a config entry."""
 
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        for mower in entry.runtime_data:
+        for mower in entry.runtime_data.mowers:
             try:
                 await mower.api.remove_device(mower.name)
             except TimeoutError:
@@ -363,7 +385,12 @@ async def async_remove_config_entry_device(
         ),
     )
     mower = next(
-        (mower for mower in config_entry.runtime_data if mower.name == mower_name), None
+        (
+            mower
+            for mower in config_entry.runtime_data.mowers
+            if mower.name == mower_name
+        ),
+        None,
     )
 
     return not bool(mower)

custom_components/mammotion/binary_sensor.py

@@ -43,7 +43,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Mammotion sensor entity."""
-    mammotion_devices = entry.runtime_data
+    mammotion_devices = entry.runtime_data.mowers
 
     for mower in mammotion_devices:
         async_add_entities(

custom_components/mammotion/button.py

@@ -76,7 +76,10 @@ class MammotionTaskButtonSensorEntityDescription(ButtonEntityDescription):
         press_fn=lambda coordinator: coordinator.join_webrtc_channel(),
         entity_category=EntityCategory.CONFIG,
     ),
-    # TODO add delete and set charging station
+    MammotionButtonSensorEntityDescription(
+        key="relocate_charging_station",
+        press_fn=lambda coordinator: coordinator.async_relocate_charging_station(),
+    ),
     # delete_charge_point
 )
 
@@ -87,7 +90,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Mammotion button sensor entity."""
-    mammotion_devices = entry.runtime_data
+    mammotion_devices = entry.runtime_data.mowers
 
     for mower in mammotion_devices:
         added_tasks: set[int] = set()

custom_components/mammotion/camera.py

@@ -55,7 +55,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Mammotion camera entities."""
-    mowers = entry.runtime_data
+    mowers = entry.runtime_data.mowers
     entities = []
     for mower in mowers:
         if not DeviceType.is_luba1(mower.device.deviceName):
@@ -153,7 +153,11 @@ def _get_mower_by_entity_id(entity_id: str):
         state = hass.states.get(entity_id)
         name = state.attributes.get("model_name")
         return next(
-            (mower for mower in entry.runtime_data if mower.device.deviceName == name),
+            (
+                mower
+                for mower in entry.runtime_data.mowers
+                if mower.device.deviceName == name
+            ),
             None,
         )
 

custom_components/mammotion/coordinator.py

@@ -25,10 +25,11 @@
     FailedRequestException,
     GatewayTimeoutException,
     NoConnectionException,
+    SetupException,
 )
 from pymammotion.aliyun.model.dev_by_account_response import Device
 from pymammotion.data.model import GenerateRouteInformation, HashList
-from pymammotion.data.model.device import MowerInfo, MowingDevice
+from pymammotion.data.model.device import MowerInfo, MowingDevice, RTKDevice
 from pymammotion.data.model.device_config import OperationSettings, create_path_order
 from pymammotion.data.model.report_info import Maintain
 from pymammotion.data.mqtt.event import DeviceNotificationEventParams, ThingEventMessage
@@ -37,13 +38,15 @@
 from pymammotion.http.model.camera_stream import (
     StreamSubscriptionResponse,
 )
-from pymammotion.http.model.http import ErrorInfo, Response
+from pymammotion.http.model.http import CheckDeviceVersion, ErrorInfo, Response
+from pymammotion.http.model.rtk import RTK
 from pymammotion.mammotion.commands.mammotion_command import MammotionCommand
 from pymammotion.mammotion.devices.mammotion import (
     ConnectionPreference,
     Mammotion,
     MammotionMixedDeviceManager,
 )
+from pymammotion.mammotion.devices.mammotion_cloud import MammotionCloud
 from pymammotion.proto import RptAct, RptInfoType, SystemUpdateBufMsg
 from pymammotion.utility.constant import WorkMode
 from pymammotion.utility.device_type import DeviceType
@@ -75,6 +78,7 @@
 REPORT_INTERVAL = timedelta(minutes=1)
 DEVICE_VERSION_INTERVAL = timedelta(days=1)
 MAP_INTERVAL = timedelta(minutes=30)
+RTK_INTERVAL = timedelta(hours=5)
 
 
 class MammotionBaseUpdateCoordinator[DataT](DataUpdateCoordinator[DataT]):
@@ -170,8 +174,8 @@ async def async_refresh_login(self) -> None:
 
     async def device_offline(self, device: MammotionMixedDeviceManager) -> None:
         device.state.online = False
-        if cloud := device.cloud():
-            await cloud.stop()
+        # if cloud := device.cloud():
+        #     await cloud.stop()
 
         loop = asyncio.get_running_loop()
         loop.call_later(900, lambda: asyncio.create_task(self.clear_update_failures()))
@@ -233,7 +237,7 @@ async def async_send_command(self, command: str, **kwargs: Any) -> bool | None:
             try:
                 if ble := device.ble():
                     # if we don't do this it will stay connected and no longer update over wifi
-                    ble.set_disconnect_strategy(True)
+                    ble.set_disconnect_strategy(disconnect=True)
                     await ble.queue_command(command, **kwargs)
 
                     return True
@@ -407,6 +411,23 @@ async def async_get_area_list(self) -> None:
         """Mowing area List."""
         await self.async_send_command("get_area_name_list", device_id=self.device.iotId)
 
+    async def async_relocate_charging_station(self):
+        """Reset charging station."""
+        await self.async_send_command("delete_charge_point")
+        # fetch charging location?
+        """
+        nav {
+          todev_get_commondata {
+            pver: 1
+            subCmd: 2
+            action: 6
+            type: 5
+            totalFrame: 1
+            currentFrame: 1
+          }
+        }
+        """
+
     async def send_command_and_update(self, command_str: str, **kwargs: Any) -> None:
         """Send command and update."""
         await self.async_send_command(command_str, **kwargs)
@@ -963,7 +984,6 @@ async def _async_update_data(self):
             if (
                 len(device.state.map.hashlist) == 0
                 or len(device.state.map.missing_hashlist()) > 0
-                or len(device.state.map.plan) == 0
             ):
                 await self.manager.start_map_sync(self.device_name)
 
@@ -1136,3 +1156,101 @@ async def _async_setup(self) -> None:
                 )
         except DeviceOfflineException:
             """Device is offline bluetooth has been attempted."""
+
+
+class MammotionRTKCoordinator(DataUpdateCoordinator[RTKDevice]):
+    """Mammotion DataUpdateCoordinator."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        config_entry: MammotionConfigEntry,
+        device: Device,
+        cloud: MammotionCloud,
+    ) -> None:
+        """Initialize global mammotion data updater."""
+        super().__init__(
+            hass=hass,
+            logger=LOGGER,
+            name=DOMAIN,
+            update_interval=RTK_INTERVAL,
+            config_entry=config_entry,
+        )
+        assert config_entry.unique_id
+        self.account = self.config_entry.data[CONF_ACCOUNTNAME]
+        self.password = self.config_entry.data[CONF_PASSWORD]
+        self.device: Device = device
+        self.device_name = device.deviceName
+        self.cloud: MammotionCloud = cloud
+        self.data: RTKDevice = RTKDevice(
+            name=self.device_name,
+            iot_id=self.device.iotId,
+            product_key=self.device.productKey,
+        )
+
+    async def _async_update_data(self):
+        """Update RTK data."""
+        try:
+            response = await self.cloud.cloud_client.get_device_properties(
+                self.device.iotId
+            )
+            if response.code == 200:
+                data = response.data
+                if ota_progress := data.otaProgress:
+                    self.data.update_check = CheckDeviceVersion.from_dict(
+                        ota_progress.value
+                    )
+                if network_info := data.networkInfo:
+                    network = json.loads(network_info.value)
+                    self.data.wifi_rssi = network["wifi_rssi"]
+                    self.data.wifi_sta_mac = network["wifi_sta_mac"]
+                    self.data.bt_mac = network["bt_mac"]
+                if coordinate := data.coordinate:
+                    coord_val = json.loads(coordinate.value)
+                    self.data.lat = coord_val["lat"]
+                    self.data.lon = coord_val["lon"]
+                if device_version := data.deviceVersion:
+                    self.data.device_version = device_version.value
+            self.data.online = True
+
+            ota_info = (
+                await self.cloud.cloud_client.mammotion_http.get_device_ota_firmware(
+                    [self.data.iot_id]
+                )
+            )
+            if check_versions := ota_info.data:
+                for check_version in check_versions:
+                    if check_version.device_id == self.data.iot_id:
+                        self.data.update_check = check_version
+            return self.data
+        except SetupException:
+            """Cloud IOT Gateway is not setup."""
+            return self.data
+        except DeviceOfflineException:
+            self.data.online = False
+        except GatewayTimeoutException:
+            """Gateway is timing out again."""
+        return self.data
+
+    async def _async_setup(self) -> None:
+        """Setup RTK data."""
+
+        rtk_response = await self.cloud.cloud_client.mammotion_http.get_rtk_devices()
+        if rtk_response.code == 0:
+            rtk_list = [
+                rtk
+                for rtk in rtk_response.data
+                if self.device.deviceName == rtk.device_name
+            ]
+            try:
+                rtk_device: RTK = next(iter(rtk_list))
+                self.data.lora_version = rtk_device.lora
+            except StopIteration:
+                """Failed to get RTK device."""
+                return
+
+    async def update_firmware(self, version: str) -> None:
+        """Update firmware."""
+        await self.cloud.cloud_client.mammotion_http.start_ota_upgrade(
+            self.device.iotId, version
+        )

custom_components/mammotion/device_tracker.py

@@ -22,7 +22,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the RTK tracker from config entry."""
-    mammotion_devices = entry.runtime_data
+    mammotion_devices = entry.runtime_data.mowers
 
     for mower in mammotion_devices:
         async_add_entities([MammotionTracker(mower.reporting_coordinator)])

custom_components/mammotion/diagnostics.py

@@ -18,7 +18,7 @@ async def async_get_config_entry_diagnostics(
     entry: MammotionConfigEntry,
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    mammotion_devices: list[MammotionMowerData] = entry.runtime_data
+    mammotion_devices: list[MammotionMowerData] = entry.runtime_data.mowers
     data = {}
     for device in mammotion_devices:
         data[device.name] = asdict(device.reporting_coordinator.data)