From 7fa80ec4d098feda903d9e32f7dbbe7d207f20a8 Mon Sep 17 00:00:00 2001
From: abmantis <amfcalt@gmail.com>
Date: Mon, 31 Jul 2023 00:16:21 +0100
Subject: [PATCH] Allow stopping in-progress movement

---
 idasen/__init__.py   | 81 +++++++++++++++++++++++++++++++-------------
 tests/test_idasen.py | 27 +++++++++++++++
 2 files changed, 85 insertions(+), 23 deletions(-)

diff --git a/idasen/__init__.py b/idasen/__init__.py
index 87b6fc6..1c36ac8 100644
--- a/idasen/__init__.py
+++ b/idasen/__init__.py
@@ -103,6 +103,8 @@ def __init__(
         self._logger = _DeskLoggingAdapter(
             logger=logging.getLogger(__name__), extra={"mac": self.mac}
         )
+        self._moving = False
+        self._move_task: Optional[asyncio.Task] = None
 
     async def __aenter__(self):
         await self._connect()
@@ -137,7 +139,7 @@ async def monitor(self, callback: Callable[[float], Awaitable[None]]):
 
         async def output_listener(char: BleakGATTCharacteristic, data: bytearray):
             height = _bytes_to_meters(data)
-            self._logger.info(f"Got data: {height}m")
+            self._logger.debug(f"Got data: {height}m")
 
             nonlocal previous_height
             if abs(height - previous_height) < 0.001:
@@ -176,6 +178,16 @@ def is_connected(self) -> bool:
         """
         return self._client.is_connected
 
+    @property
+    def is_moving(self) -> bool:
+        """
+        Check if the desk is currently being moved by this class.
+
+        Returns:
+            Boolean representing movement status.
+        """
+        return self._moving
+
     @property
     def mac(self) -> str:
         """Desk MAC address."""
@@ -235,31 +247,54 @@ async def move_to_target(self, target: float):
                 f"{self.MIN_HEIGHT:.3f}"
             )
 
-        previous_height = await self.get_height()
-        will_move_up = target > previous_height
-        while True:
-            height = await self.get_height()
-            difference = target - height
-            self._logger.debug(f"{target=} {height=} {difference=}")
-            if (height < previous_height and will_move_up) or (
-                height > previous_height and not will_move_up
-            ):
-                self._logger.warning(
-                    "stopped moving because desk safety feature kicked in"
-                )
-                return
-            if abs(difference) < 0.005:  # tolerance of 0.005 meters
-                self._logger.info(f"reached target of {target:.3f}")
-                await self.stop()
-                return
-            elif difference > 0:
-                await self.move_up()
-            elif difference < 0:
-                await self.move_down()
-            previous_height = height
+        if self._moving:
+            self._logger.error("Already moving")
+            return
+        self._moving = True
+
+        async def do_move():
+            previous_height = await self.get_height()
+            will_move_up = target > previous_height
+            while True:
+                height = await self.get_height()
+                difference = target - height
+                self._logger.debug(f"{target=} {height=} {difference=}")
+                if (height < previous_height and will_move_up) or (
+                    height > previous_height and not will_move_up
+                ):
+                    self._logger.warning(
+                        "stopped moving because desk safety feature kicked in"
+                    )
+                    return
+                if not self._moving:
+                    return
+
+                if abs(difference) < 0.005:  # tolerance of 0.005 meters
+                    self._logger.info(f"reached target of {target:.3f}")
+                    await self._stop()
+                    return
+                elif difference > 0:
+                    await self.move_up()
+                elif difference < 0:
+                    await self.move_down()
+                previous_height = height
+
+        self._move_task = asyncio.create_task(do_move())
+        await self._move_task
+        self._moving = False
 
     async def stop(self):
         """Stop desk movement."""
+        self._moving = False
+        if self._move_task:
+            self._logger.debug("Desk was moving, waiting for it to stop")
+            await self._move_task
+
+        await self._stop()
+
+    async def _stop(self):
+        """Send stop commands"""
+        self._logger.debug("Sending stop commands")
         await asyncio.gather(
             self._client.write_gatt_char(_UUID_COMMAND, _COMMAND_STOP, response=False),
             self._client.write_gatt_char(
diff --git a/tests/test_idasen.py b/tests/test_idasen.py
index 1f15e57..2b30ff5 100644
--- a/tests/test_idasen.py
+++ b/tests/test_idasen.py
@@ -136,6 +136,33 @@ async def test_move_to_target(desk: IdasenDesk, target: float):
     assert abs(await desk.get_height() - target) < 0.005
 
 
+async def test_move_stop():
+    desk = IdasenDesk(mac=desk_mac)
+    client = MockBleakClient()
+    desk._client = client
+    client.write_gatt_char = mock.AsyncMock()
+
+    async def write_gatt_char_mock(
+        uuid: str, command: bytearray, response: bool = False
+    ):
+        if client.write_gatt_char.call_count == 1:
+            assert desk.is_moving
+        # Force this method to behave asynchronously, otherwise it will block the
+        # eventloop
+        await asyncio.sleep(0.1)
+
+    client.write_gatt_char.side_effect = write_gatt_char_mock
+
+    async with desk:
+        move_task = asyncio.create_task(desk.move_to_target(0.7))
+        # Allow the move_task to start looping
+        await asyncio.sleep(0.1)
+
+        await desk.stop()
+        assert not desk.is_moving
+        assert move_task.done()
+
+
 @pytest.mark.parametrize(
     "raw, result",
     [