[commands] Add onlyWhile and onlyIf (#5291)

wpilibNewCommands/src/test/java/edu/wpi/first/wpilibj2/command/CommandDecoratorTest.java

@@ -57,6 +57,22 @@ class CommandDecoratorTest extends CommandTestBase {
     }
   }
 
+  @Test
+  void onlyWhileTest() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicBoolean condition = new AtomicBoolean(true);
+
+      Command command = new WaitCommand(10).onlyWhile(condition::get);
+
+      scheduler.schedule(command);
+      scheduler.run();
+      assertTrue(scheduler.isScheduled(command));
+      condition.set(false);
+      scheduler.run();
+      assertFalse(scheduler.isScheduled(command));
+    }
+  }
+
   @Test
   void ignoringDisableTest() {
     try (CommandScheduler scheduler = new CommandScheduler()) {
@@ -221,6 +237,26 @@ class CommandDecoratorTest extends CommandTestBase {
     }
   }
 
+  @Test
+  void onlyIfTest() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicBoolean onlyIfCondition = new AtomicBoolean(false);
+      AtomicBoolean hasRunCondition = new AtomicBoolean(false);
+
+      Command command =
+          new InstantCommand(() -> hasRunCondition.set(true)).onlyIf(onlyIfCondition::get);
+
+      scheduler.schedule(command);
+      scheduler.run();
+      assertFalse(hasRunCondition.get());
+
+      onlyIfCondition.set(true);
+      scheduler.schedule(command);
+      scheduler.run();
+      assertTrue(hasRunCondition.get());
+    }
+  }
+
   @Test
   void finallyDoTest() {
     try (CommandScheduler scheduler = new CommandScheduler()) {

wpilibNewCommands/src/test/native/cpp/frc2/command/CommandDecoratorTest.cpp

@@ -54,6 +54,24 @@ TEST_F(CommandDecoratorTest, Until) {
   EXPECT_FALSE(scheduler.IsScheduled(command));
 }
 
+TEST_F(CommandDecoratorTest, OnlyWhile) {
+  CommandScheduler scheduler = GetScheduler();
+
+  bool run = true;
+
+  auto command = RunCommand([] {}, {}).OnlyWhile([&run] { return run; });
+
+  scheduler.Schedule(command);
+
+  scheduler.Run();
+  EXPECT_TRUE(scheduler.IsScheduled(command));
+
+  run = false;
+
+  scheduler.Run();
+  EXPECT_FALSE(scheduler.IsScheduled(command));
+}
+
 TEST_F(CommandDecoratorTest, IgnoringDisable) {
   CommandScheduler scheduler = GetScheduler();
 
@@ -140,6 +158,27 @@ TEST_F(CommandDecoratorTest, Unless) {
   EXPECT_TRUE(hasRun);
 }
 
+TEST_F(CommandDecoratorTest, OnlyIf) {
+  CommandScheduler scheduler = GetScheduler();
+
+  bool hasRun = false;
+  bool onlyIfBool = false;
+
+  auto command =
+      InstantCommand([&hasRun] { hasRun = true; }, {}).OnlyIf([&onlyIfBool] {
+        return onlyIfBool;
+      });
+
+  scheduler.Schedule(command);
+  scheduler.Run();
+  EXPECT_FALSE(hasRun);
+
+  onlyIfBool = true;
+  scheduler.Schedule(command);
+  scheduler.Run();
+  EXPECT_TRUE(hasRun);
+}
+
 TEST_F(CommandDecoratorTest, FinallyDo) {
   CommandScheduler scheduler = GetScheduler();
   int first = 0;