[robotpy] Mirror robotpy's commands-v2 (#8369)

Project import generated by Copybara.

GitOrigin-RevId: 715c8e8372d936f447f2937aab6b1a22dc619126

commandsv2/src/test/python/test_swervecontrollercommand.py

@@ -0,0 +1,142 @@
+# Copyright (c) FIRST and other WPILib contributors.
+# Open Source Software; you can modify and/or share it under the terms of
+# the WPILib BSD license file in the root directory of this project.
+
+from typing import TYPE_CHECKING, List, Tuple
+import math
+
+from wpilib import Timer
+
+from wpimath.geometry import Pose2d, Rotation2d, Translation2d
+from wpimath.kinematics import (
+    SwerveModuleState,
+    SwerveModulePosition,
+    SwerveDrive4Kinematics,
+    SwerveDrive4Odometry,
+)
+from wpimath.controller import (
+    ProfiledPIDControllerRadians,
+    PIDController,
+    HolonomicDriveController,
+)
+from wpimath.trajectory import (
+    TrapezoidProfileRadians,
+    Trajectory,
+    TrajectoryConfig,
+    TrajectoryGenerator,
+)
+
+
+from util import *  # type: ignore
+
+if TYPE_CHECKING:
+    from .util import *
+
+import pytest
+
+import commands2
+
+
+def test_swervecontrollercommand():
+    timer = Timer()
+    angle = Rotation2d(0)
+
+    swerve_module_states = (
+        SwerveModuleState(0, Rotation2d(0)),
+        SwerveModuleState(0, Rotation2d(0)),
+        SwerveModuleState(0, Rotation2d(0)),
+        SwerveModuleState(0, Rotation2d(0)),
+    )
+
+    swerve_module_positions = (
+        SwerveModulePosition(0, Rotation2d(0)),
+        SwerveModulePosition(0, Rotation2d(0)),
+        SwerveModulePosition(0, Rotation2d(0)),
+        SwerveModulePosition(0, Rotation2d(0)),
+    )
+
+    rot_controller = ProfiledPIDControllerRadians(
+        1,
+        0,
+        0,
+        TrapezoidProfileRadians.Constraints(3 * math.pi, math.pi),
+    )
+
+    x_tolerance = 1 / 12.0
+    y_tolerance = 1 / 12.0
+    angular_tolerance = 1 / 12.0
+
+    wheel_base = 0.5
+    track_width = 0.5
+
+    kinematics = SwerveDrive4Kinematics(
+        Translation2d(wheel_base / 2, track_width / 2),
+        Translation2d(wheel_base / 2, -track_width / 2),
+        Translation2d(-wheel_base / 2, track_width / 2),
+        Translation2d(-wheel_base / 2, -track_width / 2),
+    )
+
+    odometry = SwerveDrive4Odometry(
+        kinematics,
+        Rotation2d(0),
+        swerve_module_positions,
+        Pose2d(0, 0, Rotation2d(0)),
+    )
+
+    def set_module_states(states):
+        nonlocal swerve_module_states
+        swerve_module_states = states
+
+    def get_robot_pose() -> Pose2d:
+        odometry.update(angle, swerve_module_positions)
+        return odometry.getPose()
+
+    with ManualSimTime() as sim:
+        subsystem = commands2.Subsystem()
+        waypoints: List[Pose2d] = []
+        waypoints.append(Pose2d(0, 0, Rotation2d(0)))
+        waypoints.append(Pose2d(1, 5, Rotation2d(3)))
+        config = TrajectoryConfig(8.8, 0.1)
+        trajectory = TrajectoryGenerator.generateTrajectory(waypoints, config)
+
+        end_state = trajectory.sample(trajectory.totalTime())
+
+        command = commands2.SwerveControllerCommand(
+            trajectory=trajectory,
+            pose=get_robot_pose,
+            kinematics=kinematics,
+            controller=HolonomicDriveController(
+                PIDController(0.6, 0, 0),
+                PIDController(0.6, 0, 0),
+                rot_controller,
+            ),
+            outputModuleStates=set_module_states,
+            requirements=(subsystem,),
+        )
+
+        timer.restart()
+
+        command.initialize()
+        while not command.isFinished():
+            command.execute()
+            angle = trajectory.sample(timer.get()).pose.rotation()
+
+            for i in range(0, len(swerve_module_positions)):
+                swerve_module_positions[i].distance += (
+                    swerve_module_states[i].speed * 0.005
+                )
+                swerve_module_positions[i].angle = swerve_module_states[i].angle
+
+            sim.step(0.005)
+
+        timer.stop()
+        command.end(True)
+
+        assert end_state.pose.X() == pytest.approx(get_robot_pose().X(), x_tolerance)
+
+        assert end_state.pose.Y() == pytest.approx(get_robot_pose().Y(), y_tolerance)
+
+        assert end_state.pose.rotation().radians() == pytest.approx(
+            get_robot_pose().rotation().radians(),
+            angular_tolerance,
+        )