Support rotation counter setting in EncoderMotor class

packages/pma/pitop/pma/encoder_motor.py

@@ -1,5 +1,4 @@
 import atexit
-import time
 from math import floor, pi
 
 from pitop.core.mixins import Recreatable, Stateful
@@ -69,8 +68,7 @@ def __init__(
         self.__forward_direction = forward_direction
         self.__wheel_diameter = wheel_diameter
 
-        self.__prev_time_dist_cnt = time.time()
-        self.__previous_reading_odometer = 0
+        self.__rotation_counter_offset = 0  # Tracks rotation counter
 
         atexit.register(self.stop)
 
@@ -286,15 +284,40 @@ def rotation_counter(self):
         the output shaft.
         """
 
-        dc_motor_rotation_counter = (
-            self.__motor_core.odometer() * self.__forward_direction
+        dc_motor_rotation_counter = self.__motor_core.odometer()
+        adjusted_dc_motor_rotation_counter = dc_motor_rotation_counter - (
+            self.__rotation_counter_offset
         )
         output_shaft_rotation_counter = round(
-            dc_motor_rotation_counter / self.MMK_STANDARD_GEAR_RATIO, 1
+            adjusted_dc_motor_rotation_counter / self.MMK_STANDARD_GEAR_RATIO, 1
         )
 
         return output_shaft_rotation_counter
 
+    def reset_rotation_counter(self, rotations=0):
+        """Set the rotation counter to a specific rotations value.
+
+        This method sets the rotation counter offset to the difference between the
+        current odometer value and the provided rotations value, setting the current
+        rotations counter but having it still change with the odometer.
+
+        We use an offset to track the difference between the current odometer value and
+        the user provided rotations value, so that we can report the current rotations
+        relative to this fixed value as they change.
+
+        :type rotations: int or float
+        :param rotations:
+            New number of rotations to set the rotation counter to
+        """
+        if not isinstance(rotations, (float, int)):
+            raise ValueError("Rotation counter must be a number")
+
+        current_odometer = self.__motor_core.odometer()
+        # The offset is now set so that rotation_counter returns the provided value for the current odometer reading
+        self.__rotation_counter_offset = (
+            current_odometer - rotations * self.MMK_STANDARD_GEAR_RATIO
+        )
+
     @property
     def torque_limited(self):
         """Check if the actual motor speed or RPM does not match the target
@@ -449,7 +472,9 @@ def distance(self):
         value being set.
         """
 
-        return self.wheel_circumference * self.rotation_counter
+        return (
+            self.wheel_circumference * self.rotation_counter * self.__forward_direction
+        )
 
     @property
     def max_speed(self):

tests/test_pma_encoder_motor.py

@@ -141,3 +141,63 @@ def test_set_target_speed_fails_when_requesting_an_out_of_range_speed(
                 set_target_rpm_mock.assert_called_with(
                     target_speed_in_rpm, direction, target_motor_rotations
                 )
+
+    @patch("pitop.pma.encoder_motor.EncoderMotorController.odometer")
+    def test_reset_rotation_counter_method(self, mock_odometer):
+        """Test that setting rotation_counter sets the offset correctly."""
+        # Mock odometer call
+        mock_odometer.return_value = 456
+
+        encoder_motor = EncoderMotor(
+            port_name="M1",
+            forward_direction=ForwardDirection.CLOCKWISE,
+            braking_type=BrakingType.COAST,
+        )
+
+        # Initial value is based on the odometer
+        expected_rotations = round((456 / encoder_motor.MMK_STANDARD_GEAR_RATIO), 1)
+        self.assertEqual(encoder_motor.rotation_counter, expected_rotations)
+
+        # Set rotation_counter property to a few values; the property will now return the same value
+        encoder_motor.reset_rotation_counter()
+        self.assertEqual(encoder_motor.rotation_counter, 0)
+
+        encoder_motor.reset_rotation_counter(100)
+        self.assertEqual(encoder_motor.rotation_counter, 100)
+
+        # When the EncoderMotor moves, the odometer returns a new value, which causes the rotation counter to be updated
+        mock_odometer.return_value = 1500.0
+        # Calculate expected_rotations based on the new offset, based on the current and previous odometer readings
+        expected_rotations = round(
+            100 + (1500.0 - 456) / encoder_motor.MMK_STANDARD_GEAR_RATIO, 1
+        )
+        self.assertEqual(encoder_motor.rotation_counter, expected_rotations)
+
+        # Make sure this still works when the odometer moves in the opposite direction and returns a negative value
+        mock_odometer.return_value = -1000.0
+        expected_rotations = round(
+            expected_rotations
+            + (-1000.0 - 1500) / encoder_motor.MMK_STANDARD_GEAR_RATIO,
+            1,
+        )
+        self.assertEqual(encoder_motor.rotation_counter, expected_rotations)
+
+        # Setter also handles negative values
+        encoder_motor.reset_rotation_counter(-100)
+        self.assertEqual(encoder_motor.rotation_counter, -100)
+
+    @patch("pitop.pma.encoder_motor.EncoderMotorController.odometer")
+    def test_reset_rotation_counter_error_handling(self, mock_odometer):
+        """Test that setting rotation_counter setter raises an error when given an invalid value."""
+        # Mock odometer call
+        mock_odometer.return_value = 0
+
+        encoder_motor = EncoderMotor(
+            port_name="M1",
+            forward_direction=ForwardDirection.CLOCKWISE,
+            braking_type=BrakingType.COAST,
+        )
+
+        for invalid_value in ("100", b"123"):
+            with self.assertRaises(ValueError):
+                encoder_motor.reset_rotation_counter(invalid_value)