Add an option in @InitWPILib to change the period

testing/src/main/java/com/team2813/lib2813/testing/junit/jupiter/InitWPILib.java

@@ -15,6 +15,7 @@
 */
 package com.team2813.lib2813.testing.junit.jupiter;
 
+import edu.wpi.first.wpilibj.TimedRobot;
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -42,7 +43,7 @@
  *   @Test
  *   public void takesFlight(CommandTester commandTester) {
  *     var flight = new FlightSubsystem();
- *     Command takeOff = flight.createTakeOffCommandCommand();
+ *     Command takeOff = flight.createTakeOffCommand();
  *
  *     commandTester.runUntilComplete(takeOff);
  *
@@ -56,4 +57,10 @@
 @Target(ElementType.TYPE)
 @Retention(RetentionPolicy.RUNTIME)
 @ExtendWith(InitWPILibExtension.class)
-public @interface InitWPILib {}
+public @interface InitWPILib {
+  /**
+   * The time stepped in between each periodic run when running a command. This is equivalent to the
+   * {@code period} argument in {@link TimedRobot#TimedRobot(double)}}.
+   */
+  double periodicPeriod() default TimedRobot.kDefaultPeriod;
+}

testing/src/main/java/com/team2813/lib2813/testing/junit/jupiter/InitWPILibExtension.java

@@ -17,7 +17,6 @@
 
 import edu.wpi.first.hal.HAL;
 import edu.wpi.first.wpilibj.RuntimeType;
-import edu.wpi.first.wpilibj.TimedRobot;
 import edu.wpi.first.wpilibj.simulation.DriverStationSim;
 import edu.wpi.first.wpilibj.simulation.SimHooks;
 import edu.wpi.first.wpilibj2.command.CommandScheduler;
@@ -26,9 +25,12 @@
 import org.junit.jupiter.api.extension.BeforeAllCallback;
 import org.junit.jupiter.api.extension.Extension;
 import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.api.extension.ExtensionContext.Namespace;
+import org.junit.jupiter.api.extension.ExtensionContext.Store;
 import org.junit.jupiter.api.extension.ParameterContext;
 import org.junit.jupiter.api.extension.ParameterResolutionException;
 import org.junit.jupiter.api.extension.ParameterResolver;
+import org.junit.platform.commons.support.AnnotationSupport;
 
 /** JUnit Jupiter extension for testing code that depends on WPILib. */
 final class InitWPILibExtension
@@ -38,9 +40,13 @@ final class InitWPILibExtension
         BeforeAllCallback,
         ParameterResolver {
   private static final double NANOS_PER_SECOND = 1_000_000_000d;
+  private static final Namespace NAMESPACE = Namespace.create(InitWPILibExtension.class);
+  private static final StoreKey<InitWPILib> ANNOTATION_KEY = StoreKey.of(InitWPILib.class);
 
   @Override
   public void beforeAll(ExtensionContext context) {
+    Store store = context.getStore(NAMESPACE);
+    ANNOTATION_KEY.put(store, getAnnotation(context));
     // Ensure the Hardware Abstraction Layer is initialized before we try to use it. This logic is
     // based on a comment from Peter Johnson at
     // https://www.chiefdelphi.com/t/driverstation-getalliance-in-gradle-test/
@@ -85,18 +91,28 @@ public boolean supportsParameter(
   public CommandTester resolveParameter(
       ParameterContext parameterContext, ExtensionContext extensionContext) {
     CommandScheduler scheduler = CommandScheduler.getInstance();
+    InitWPILib annotation = ANNOTATION_KEY.get(extensionContext.getStore(NAMESPACE));
 
     return command -> {
       SimHooks.pauseTiming();
       try {
         scheduler.schedule(command);
         do {
           scheduler.run();
-          SimHooks.stepTiming(TimedRobot.kDefaultPeriod);
+          SimHooks.stepTiming(annotation.periodicPeriod());
         } while (scheduler.isScheduled(command));
       } finally {
         SimHooks.resumeTiming();
       }
     };
   }
+
+  private static InitWPILib getAnnotation(ExtensionContext context) {
+    return AnnotationSupport.findAnnotation(
+            context.getRequiredTestClass(), InitWPILib.class, context.getEnclosingTestClasses())
+        .orElseThrow(
+            () ->
+                new IllegalStateException(
+                    "Could not find enclosed class annotated with @InitWPILib"));
+  }
 }

testing/src/test/java/com/team2813/lib2813/testing/junit/jupiter/InitWPILibExtensionTest.java

@@ -23,6 +23,8 @@
 import edu.wpi.first.hal.HAL;
 import edu.wpi.first.wpilibj.DriverStation;
 import edu.wpi.first.wpilibj.RobotState;
+import edu.wpi.first.wpilibj.TimedRobot;
+import edu.wpi.first.wpilibj.Timer;
 import edu.wpi.first.wpilibj.simulation.DriverStationSim;
 import edu.wpi.first.wpilibj2.command.Command;
 import edu.wpi.first.wpilibj2.command.CommandScheduler;
@@ -81,6 +83,7 @@ public void verifyFakeCommandNotScheduledBeforeTest() {
     @Test
     @Order(2)
     public void verifyFakeCommandNotScheduledAfterTest(CommandTester commandTester) {
+      assertThat(commandTester).isNotNull();
       CommandScheduler commandScheduler = CommandScheduler.getInstance();
       assertWithMessage("Expect all commands to have been cancelled")
           .that(commandScheduler.isScheduled(FAKE_COMMAND))
@@ -143,6 +146,52 @@ private void withDriverStationTemporarilyEnabled(Runnable runnable) {
     }
   }
 
+  @InitWPILib(periodicPeriod = 0.05)
+  @Tag("ignore-outside-testkit")
+  public static class PeriodicDurationHalfTest {
+    @Test
+    public void periodicPeriodIsCorrectTest(CommandTester commandTester) {
+      PeriodicElapsedCommand command = new PeriodicElapsedCommand();
+      commandTester.runUntilComplete(command);
+      assertThat(command.executionTime()).isWithin(0.1).of(0.05);
+    }
+  }
+
+  @InitWPILib(periodicPeriod = 0.08)
+  @Tag("ignore-outside-testkit")
+  public static class PeriodicDurationFourFifthsTest {
+    @Test
+    public void periodicPeriodIsCorrectTest(CommandTester commandTester) {
+      PeriodicElapsedCommand command = new PeriodicElapsedCommand();
+      commandTester.runUntilComplete(command);
+      assertThat(command.executionTime()).isWithin(0.1).of(0.08);
+    }
+  }
+
+  @InitWPILib
+  @Tag("ignore-outside-testkit")
+  public static class PeriodicDurationDefaultTest {
+    @Test
+    public void periodicPeriodIsCorrectTest(CommandTester commandTester) {
+      PeriodicElapsedCommand command = new PeriodicElapsedCommand();
+      commandTester.runUntilComplete(command);
+      assertThat(command.executionTime()).isWithin(0.1).of(TimedRobot.kDefaultPeriod);
+    }
+  }
+
+  @Test
+  void verifyPeriodicPeriodConfig() {
+    EngineExecutionResults results =
+        EngineTestKit.engine("junit-jupiter")
+            .selectors(
+                selectClass(PeriodicDurationHalfTest.class),
+                selectClass(PeriodicDurationFourFifthsTest.class),
+                selectClass(PeriodicDurationDefaultTest.class))
+            .execute();
+
+    assertHasNoFailures(results);
+  }
+
   private static class VerifiableCommand extends Command {
     private static final int EXPECTED_EXECUTION_COUNT = 4;
     private int initializedCount = 0;
@@ -173,4 +222,47 @@ public boolean isFinished() {
       return executionCount >= EXPECTED_EXECUTION_COUNT;
     }
   }
+
+  private static class PeriodicElapsedCommand extends Command {
+    private boolean completed;
+    private int executeCount;
+    private double ex1Time;
+    private double ex2Time;
+
+    @Override
+    public void initialize() {
+      executeCount = 0;
+      completed = false;
+    }
+
+    @Override
+    public void execute() {
+      executeCount++;
+      assertWithMessage("execute() must only be called twice!").that(executeCount).isLessThan(3);
+
+      double now = Timer.getFPGATimestamp();
+      if (executeCount == 1) {
+        ex1Time = now;
+      } else {
+        ex2Time = now;
+      }
+    }
+
+    @Override
+    public boolean isFinished() {
+      return executeCount == 2;
+    }
+
+    @Override
+    public void end(boolean interrupted) {
+      completed = !interrupted;
+    }
+
+    public double executionTime() {
+      assertWithMessage("This command must run to completion before getting the execution time!")
+          .that(completed)
+          .isTrue();
+      return ex2Time - ex1Time;
+    }
+  }
 }