From e9f61e5a3782422ad50f54b56b4a0bd7022dae6c Mon Sep 17 00:00:00 2001
From: Robert Resch <robert@resch.dev>
Date: Mon, 13 Dec 2021 17:24:34 +0100
Subject: [PATCH] improve charging detection (#47)

---
 deebot_client/commands/charge.py       | 18 ++++++++++++++----
 deebot_client/commands/charge_state.py |  5 +++--
 deebot_client/commands/common.py       |  5 ++---
 deebot_client/commands/const.py        |  3 +++
 deebot_client/vacuum_bot.py            | 21 +++++++++++++++++++++
 5 files changed, 43 insertions(+), 9 deletions(-)
 create mode 100644 deebot_client/commands/const.py

diff --git a/deebot_client/commands/charge.py b/deebot_client/commands/charge.py
index 203145b1..f1d43c27 100644
--- a/deebot_client/commands/charge.py
+++ b/deebot_client/commands/charge.py
@@ -2,9 +2,13 @@
 from typing import Any, Dict
 
 from ..events import StatusEvent
-from ..message import HandlingResult, HandlingState
+from ..logging_filter import get_logger
+from ..message import HandlingResult
 from ..models import VacuumState
 from .common import EventBus, _ExecuteCommand
+from .const import CODE
+
+_LOGGER = get_logger(__name__)
 
 
 class Charge(_ExecuteCommand):
@@ -21,8 +25,14 @@ def _handle_body(cls, event_bus: EventBus, body: Dict[str, Any]) -> HandlingResu
 
         :return: A message response
         """
-        response = super()._handle_body(event_bus, body)
-        if response.state == HandlingState.SUCCESS:
+        code = int(body.get(CODE, -1))
+        if code == 0:
             event_bus.notify(StatusEvent(True, VacuumState.RETURNING))
+            return HandlingResult.success()
+
+        if code == 30007:
+            # bot is already charging
+            event_bus.notify(StatusEvent(True, VacuumState.DOCKED))
+            return HandlingResult.success()
 
-        return response
+        return super()._handle_body(event_bus, body)
diff --git a/deebot_client/commands/charge_state.py b/deebot_client/commands/charge_state.py
index 98d62afa..f0f0a277 100644
--- a/deebot_client/commands/charge_state.py
+++ b/deebot_client/commands/charge_state.py
@@ -4,7 +4,8 @@
 from ..events import StatusEvent
 from ..message import HandlingResult
 from ..models import VacuumState
-from .common import _CODE, EventBus, _NoArgsCommand
+from .common import EventBus, _NoArgsCommand
+from .const import CODE
 
 
 class GetChargeState(_NoArgsCommand):
@@ -26,7 +27,7 @@ def _handle_body_data_dict(
 
     @classmethod
     def _handle_body(cls, event_bus: EventBus, body: Dict[str, Any]) -> HandlingResult:
-        if body.get(_CODE, 0) == 0:
+        if body.get(CODE, 0) == 0:
             # Call this also if code is not in the body
             return super()._handle_body(event_bus, body)
 
diff --git a/deebot_client/commands/common.py b/deebot_client/commands/common.py
index 75a1ca53..7d1a590b 100644
--- a/deebot_client/commands/common.py
+++ b/deebot_client/commands/common.py
@@ -8,11 +8,10 @@
 from ..events.event_bus import EventBus
 from ..logging_filter import get_logger
 from ..message import HandlingResult, HandlingState, Message
+from .const import CODE
 
 _LOGGER = get_logger(__name__)
 
-_CODE = "code"
-
 
 @dataclass(frozen=True)
 class CommandResult(HandlingResult):
@@ -87,7 +86,7 @@ def _handle_body(cls, event_bus: EventBus, body: Dict[str, Any]) -> HandlingResu
         :return: A message response
         """
         # Success event looks like { "code": 0, "msg": "ok" }
-        if body.get(_CODE, -1) == 0:
+        if body.get(CODE, -1) == 0:
             return HandlingResult.success()
 
         _LOGGER.warning('Command "%s" was not successfully. body=%s', cls.name, body)
diff --git a/deebot_client/commands/const.py b/deebot_client/commands/const.py
new file mode 100644
index 00000000..0f2dcc0f
--- /dev/null
+++ b/deebot_client/commands/const.py
@@ -0,0 +1,3 @@
+"""Command constants module."""
+
+CODE = "code"
diff --git a/deebot_client/vacuum_bot.py b/deebot_client/vacuum_bot.py
index 5ebfc9af..39ef5c35 100644
--- a/deebot_client/vacuum_bot.py
+++ b/deebot_client/vacuum_bot.py
@@ -14,6 +14,8 @@
 from .events import (
     CleanLogEvent,
     LifeSpanEvent,
+    PositionsEvent,
+    PositionType,
     StatsEvent,
     StatusEvent,
     TotalStatsEvent,
@@ -52,6 +54,25 @@ def __init__(
 
         self.map: Final[Map] = Map(self.execute_command, self.events)
 
+        async def on_pos(event: PositionsEvent) -> None:
+            if self._status == StatusEvent(True, VacuumState.DOCKED):
+                return
+
+            deebot = next(p for p in event.positions if p.type == PositionType.DEEBOT)
+
+            if deebot:
+                on_charger = filter(
+                    lambda p: p.type == PositionType.CHARGER
+                    and p.x == deebot.x
+                    and p.y == deebot.y,
+                    event.positions,
+                )
+                if on_charger:
+                    # deebot on charger so the status should be docked... Checking
+                    self.events.request_refresh(StatusEvent)
+
+        self.events.subscribe(PositionsEvent, on_pos)
+
         async def on_status(event: StatusEvent) -> None:
             last_status = self._status
             self._status = event