From 025bc543498c1c4a74827772fa99a129f786437f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stanis=C5=82aw=20Malinowski?=
 <56644812+stan-dot@users.noreply.github.com>
Date: Tue, 27 Aug 2024 11:26:42 +0100
Subject: [PATCH] make motor locatable (#541)

* make motor locatable

* trying to get the test right

* update the test
---
 src/ophyd_async/epics/motor.py | 17 +++++++++++++++--
 tests/epics/test_motor.py      | 15 +++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/ophyd_async/epics/motor.py b/src/ophyd_async/epics/motor.py
index 8c20d238fe..309e62033f 100644
--- a/src/ophyd_async/epics/motor.py
+++ b/src/ophyd_async/epics/motor.py
@@ -1,7 +1,13 @@
 import asyncio
 from typing import Optional
 
-from bluesky.protocols import Flyable, Movable, Preparable, Stoppable
+from bluesky.protocols import (
+    Flyable,
+    Locatable,
+    Location,
+    Preparable,
+    Stoppable,
+)
 from pydantic import BaseModel, Field
 
 from ophyd_async.core import (
@@ -51,7 +57,7 @@ class FlyMotorInfo(BaseModel):
     timeout: CalculatableTimeout = Field(frozen=True, default=CalculateTimeout)
 
 
-class Motor(StandardReadable, Movable, Stoppable, Flyable, Preparable):
+class Motor(StandardReadable, Locatable, Stoppable, Flyable, Preparable):
     """Device that moves a motor record"""
 
     def __init__(self, prefix: str, name="") -> None:
@@ -193,6 +199,13 @@ async def _prepare_velocity(
         await self.velocity.set(abs(max_speed))
         return fly_velocity
 
+    async def locate(self) -> Location[float]:
+        location: Location = {
+            "setpoint": await self.user_setpoint.get_value(),
+            "readback": await self.user_readback.get_value(),
+        }
+        return location
+
     async def _prepare_motor_path(
         self, fly_velocity: float, start_position: float, end_position: float
     ) -> float:
diff --git a/tests/epics/test_motor.py b/tests/epics/test_motor.py
index d83a90b3f4..460eebb092 100644
--- a/tests/epics/test_motor.py
+++ b/tests/epics/test_motor.py
@@ -13,6 +13,7 @@
     MockSignalBackend,
     SignalRW,
     callback_on_mock_put,
+    mock_puts_blocked,
     observe_value,
     set_mock_put_proceeds,
     set_mock_value,
@@ -317,3 +318,17 @@ async def test_complete(sim_motor: motor.Motor) -> None:
     assert not sim_motor._fly_status.done
     await sim_motor.complete()
     assert sim_motor._fly_status.done
+
+
+async def test_locatable(sim_motor: motor.Motor) -> None:
+    callback_on_mock_put(
+        sim_motor.user_setpoint,
+        lambda x, *_, **__: set_mock_value(sim_motor.user_readback, x),
+    )
+    assert (await sim_motor.locate())["readback"] == 0
+    async with mock_puts_blocked(sim_motor.user_setpoint):
+        move_status = sim_motor.set(10)
+        assert (await sim_motor.locate())["readback"] == 0
+    await move_status
+    assert (await sim_motor.locate())["readback"] == 10
+    assert (await sim_motor.locate())["setpoint"] == 10