From 68d009e7c4d67b982ad4bb7dc5e8564f6da60526 Mon Sep 17 00:00:00 2001 From: Jorge Capona Date: Thu, 24 Jul 2025 16:22:22 -0400 Subject: [PATCH] Enhance DriveController with additional movement parameters Added optional parameters for wheel rotations and speed in meters per second to the forward and backward methods. --- .../pitop/robotics/drive_controller.py | 54 ++++++- tests/test_drive_controller.py | 142 +++++++++++++++++- 2 files changed, 189 insertions(+), 7 deletions(-) diff --git a/packages/robotics/pitop/robotics/drive_controller.py b/packages/robotics/pitop/robotics/drive_controller.py index 9fc426b5e..1bb0fa9ba 100644 --- a/packages/robotics/pitop/robotics/drive_controller.py +++ b/packages/robotics/pitop/robotics/drive_controller.py @@ -129,11 +129,21 @@ def robot_move( ) self._set_motor_speeds(speed_left, speed_right, distance) + def _rotations_to_distance(self, wheel_rotations: Union[int, float]) -> float: + return wheel_rotations * self.left_motor.wheel_circumference + + def _speed_meters_per_second_to_factor( + self, speed_meters_per_second: Union[int, float] + ) -> float: + return speed_meters_per_second / self.max_motor_speed + def forward( self, - speed_factor: Union[int, float], + speed_factor: Optional[Union[int, float]] = None, hold: bool = False, distance: Optional[Union[int, float]] = None, + wheel_rotations: Optional[Union[int, float]] = None, + speed_meters_per_second: Optional[Union[int, float]] = None, ) -> None: """Move the robot forward. @@ -145,19 +155,43 @@ def forward( subsequent movements to use the speed set as the base speed. :param float distance: Distance to travel in meters. If not provided, the robot will move indefinitely + :param float wheel_rotations: Number of wheel rotations to travel. Can't be set at the same time as distance. + :param float speed_meters_per_second: Speed in meters per second. Can't be set at the same time as speed_factor. """ - linear_speed_x = self.max_motor_speed * speed_factor + + if (speed_factor is None and speed_meters_per_second is None) or ( + speed_factor is not None and speed_meters_per_second is not None + ): + raise ValueError( + "Either speed_factor or speed_meters_per_second must be provided, but not both" + ) + + if distance is not None and wheel_rotations is not None: + raise ValueError( + "Either distance or wheel_rotations must be provided, but not both" + ) + + linear_speed_x = speed_meters_per_second + if speed_factor: + linear_speed_x = self.max_motor_speed * speed_factor + if hold: self._linear_speed_x_hold = linear_speed_x else: self._linear_speed_x_hold = 0 + + if wheel_rotations: + distance = self._rotations_to_distance(wheel_rotations) + self.robot_move(linear_speed=linear_speed_x, angular_speed=0, distance=distance) def backward( self, - speed_factor: Union[int, float], + speed_factor: Optional[Union[int, float]] = None, hold: bool = False, distance: Optional[Union[int, float]] = None, + wheel_rotations: Optional[Union[int, float]] = None, + speed_meters_per_second: Optional[Union[int, float]] = None, ) -> None: """Move the robot backward. @@ -168,8 +202,16 @@ def backward( subsequent movements to use the speed set as the base speed. :param float distance: Distance to travel in meters. If not provided, the robot will move indefinitely + :param float wheel_rotations: Number of wheel rotations to travel. If not + provided, the robot will move indefinitely """ - self.forward(-speed_factor, hold, distance) + self.forward( + -speed_factor if speed_factor else None, + hold, + distance, + wheel_rotations, + -speed_meters_per_second if speed_meters_per_second else None, + ) def left( self, @@ -300,8 +342,8 @@ def _set_motor_speeds( ) -> None: """Set the speed of the left and right motors. - :param float left_speed: Speed for the left motor. - :param float right_speed: Speed for the right motor. + :param float left_speed: Speed for the left motor in meters per second. + :param float right_speed: Speed for the right motor in meters per second. :param float distance: Distance to travel in meters. """ if distance is None: diff --git a/tests/test_drive_controller.py b/tests/test_drive_controller.py index 85de6747d..cbe34fa7c 100644 --- a/tests/test_drive_controller.py +++ b/tests/test_drive_controller.py @@ -74,7 +74,7 @@ def test_backward_calls_forward_method(): with patch.object(d, "forward") as forward_mock: d.backward(speed_factor, hold=hold) - forward_mock.assert_called_once_with(-speed_factor, hold, distance) + forward_mock.assert_called_once_with(-speed_factor, hold, distance, None, None) def test_stop_sets_angular_and_linear_speeds_to_zero(): @@ -177,3 +177,143 @@ def test_rotate_speeds(): sleep_mock.assert_called_once_with(0.4) sleep_mock.reset_mock() set_motor_speeds_mock.reset_mock() + + +def test_forward_raises_exception_on_invalid_speed_parameter_combinations(): + """'forward' method raises ValueError when invalid speed parameter combinations are used.""" + d = DriveController() + + # Neither speed_factor nor speed_meters_per_second provided + with pytest.raises( + ValueError, + match="Either speed_factor or speed_meters_per_second must be provided, but not both", + ): + d.forward() + + # Both speed_factor and speed_meters_per_second provided also raises exception + with pytest.raises( + ValueError, + match="Either speed_factor or speed_meters_per_second must be provided, but not both", + ): + d.forward(speed_factor=0.5, speed_meters_per_second=1.0) + + +def test_forward_raises_exception_on_invalid_distance_parameter_combinations(): + """'forward' method raises ValueError when invalid distance parameter combinations are used.""" + d = DriveController() + + # Providing both distance and wheel_rotations raises exception + with pytest.raises( + ValueError, + match="Either distance or wheel_rotations must be provided, but not both", + ): + d.forward(speed_factor=0.5, distance=1.0, wheel_rotations=2.0) + + +def test_backward_raises_exception_on_invalid_speed_parameter_combinations(): + """'backward' method raises ValueError when invalid speed parameter combinations are used.""" + d = DriveController() + + # Neither speed_factor nor speed_meters_per_second provided + with pytest.raises( + ValueError, + match="Either speed_factor or speed_meters_per_second must be provided, but not both", + ): + d.backward() + + # Both speed_factor and speed_meters_per_second provided also raises exception + with pytest.raises( + ValueError, + match="Either speed_factor or speed_meters_per_second must be provided, but not both", + ): + d.backward(speed_factor=0.5, speed_meters_per_second=1.0) + + +def test_backward_raises_exception_on_invalid_distance_parameter_combinations(): + """'backward' method raises ValueError when invalid distance parameter combinations are used.""" + d = DriveController() + + # Providing both distance and wheel_rotations raises exception + with pytest.raises( + ValueError, + match="Either distance or wheel_rotations must be provided, but not both", + ): + d.backward(speed_factor=0.5, distance=1.0, wheel_rotations=2.0) + + +@pytest.mark.parametrize( + "method_name, speed_factor, speed_meters_per_second, distance, wheel_rotations", + [ + # Test forward with various invalid combinations + ("forward", None, None, None, None), # No speed parameters + ("forward", 0.5, 1.0, None, None), # Both speed parameters + ("forward", 0.5, None, 1.0, 2.0), # Both distance parameters + ("forward", None, None, 1.0, 2.0), # No speed + both distance parameters + ("forward", 0.5, 1.0, 1.0, 2.0), # All invalid combinations + # Test backward with various invalid combinations + ("backward", None, None, None, None), # No speed parameters + ("backward", 0.5, 1.0, None, None), # Both speed parameters + ("backward", 0.5, None, 1.0, 2.0), # Both distance parameters + ("backward", None, None, 1.0, 2.0), # No speed + both distance parameters + ("backward", 0.5, 1.0, 1.0, 2.0), # All invalid combinations + ], +) +def test_forward_backward_parameter_validation( + method_name, speed_factor, speed_meters_per_second, distance, wheel_rotations +): + """Parametric test for forward/backward parameter validation.""" + d = DriveController() + method = getattr(d, method_name) + + with pytest.raises(ValueError): + method( + speed_factor=speed_factor, + speed_meters_per_second=speed_meters_per_second, + distance=distance, + wheel_rotations=wheel_rotations, + ) + + +def test_forward_backward_valid_parameter_combinations(): + """Test that valid parameter combinations work correctly for forward/backward.""" + d = DriveController() + + # Valid combinations that should not raise exceptions + with patch.object(d, "robot_move") as robot_move_mock: + # Valid: speed_factor only + d.forward(speed_factor=0.5) + robot_move_mock.assert_called() + robot_move_mock.reset_mock() + + # Valid: speed_meters_per_second only + d.forward(speed_meters_per_second=1.0) + robot_move_mock.assert_called() + robot_move_mock.reset_mock() + + # Valid: speed_factor + distance + d.forward(speed_factor=0.5, distance=1.0) + robot_move_mock.assert_called() + robot_move_mock.reset_mock() + + # Valid: speed_factor + wheel_rotations + d.forward(speed_factor=0.5, wheel_rotations=2.0) + robot_move_mock.assert_called() + robot_move_mock.reset_mock() + + # Valid: speed_meters_per_second + distance + d.forward(speed_meters_per_second=1.0, distance=1.0) + robot_move_mock.assert_called() + robot_move_mock.reset_mock() + + # Valid: speed_meters_per_second + wheel_rotations + d.forward(speed_meters_per_second=1.0, wheel_rotations=2.0) + robot_move_mock.assert_called() + robot_move_mock.reset_mock() + + # Test same combinations for backward + d.backward(speed_factor=0.5) + robot_move_mock.assert_called() + robot_move_mock.reset_mock() + + d.backward(speed_meters_per_second=1.0, distance=1.0) + robot_move_mock.assert_called()