[commands] Add FinallyDo and HandleInterrupt decorators (#4412)

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

@@ -4,12 +4,14 @@
 
 package edu.wpi.first.wpilibj2.command;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import edu.wpi.first.hal.HAL;
 import edu.wpi.first.wpilibj.simulation.SimHooks;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.parallel.ResourceLock;
 
@@ -234,4 +236,75 @@ class CommandDecoratorTest extends CommandTestBase {
       assertTrue(hasRunCondition.get());
     }
   }
+
+  @Test
+  void finallyDoTest() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicInteger first = new AtomicInteger(0);
+      AtomicInteger second = new AtomicInteger(0);
+
+      Command command =
+          new FunctionalCommand(
+                  () -> {},
+                  () -> {},
+                  interrupted -> {
+                    if (!interrupted) {
+                      first.incrementAndGet();
+                    }
+                  },
+                  () -> true)
+              .finallyDo(
+                  interrupted -> {
+                    if (!interrupted) {
+                      // to differentiate between "didn't run" and "ran before command's `end()`
+                      second.addAndGet(1 + first.get());
+                    }
+                  });
+
+      scheduler.schedule(command);
+      assertEquals(0, first.get());
+      assertEquals(0, second.get());
+      scheduler.run();
+      assertEquals(1, first.get());
+      // if `second == 0`, neither of the lambdas ran.
+      // if `second == 1`, the second lambda ran before the first one
+      assertEquals(2, second.get());
+    }
+  }
+
+  // handleInterruptTest() implicitly tests the interrupt=true branch of finallyDo()
+  @Test
+  void handleInterruptTest() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicInteger first = new AtomicInteger(0);
+      AtomicInteger second = new AtomicInteger(0);
+
+      Command command =
+          new FunctionalCommand(
+                  () -> {},
+                  () -> {},
+                  interrupted -> {
+                    if (interrupted) {
+                      first.incrementAndGet();
+                    }
+                  },
+                  () -> false)
+              .handleInterrupt(
+                  () -> {
+                    // to differentiate between "didn't run" and "ran before command's `end()`
+                    second.addAndGet(1 + first.get());
+                  });
+
+      scheduler.schedule(command);
+      scheduler.run();
+      assertEquals(0, first.get());
+      assertEquals(0, second.get());
+
+      scheduler.cancel(command);
+      assertEquals(1, first.get());
+      // if `second == 0`, neither of the lambdas ran.
+      // if `second == 1`, the second lambda ran before the first one
+      assertEquals(2, second.get());
+    }
+  }
 }

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

@@ -7,6 +7,7 @@
 #include "CommandTestBase.h"
 #include "frc2/command/ConditionalCommand.h"
 #include "frc2/command/EndlessCommand.h"
+#include "frc2/command/FunctionalCommand.h"
 #include "frc2/command/InstantCommand.h"
 #include "frc2/command/ParallelRaceGroup.h"
 #include "frc2/command/PerpetualCommand.h"
@@ -152,3 +153,62 @@ TEST_F(CommandDecoratorTest, Unless) {
   scheduler.Run();
   EXPECT_TRUE(hasRun);
 }
+
+TEST_F(CommandDecoratorTest, FinallyDo) {
+  CommandScheduler scheduler = GetScheduler();
+  int first = 0;
+  int second = 0;
+  CommandPtr command = FunctionalCommand([] {}, [] {},
+                                         [&first](bool interrupted) {
+                                           if (!interrupted) {
+                                             first++;
+                                           }
+                                         },
+                                         [] { return true; })
+                           .FinallyDo([&first, &second](bool interrupted) {
+                             if (!interrupted) {
+                               // to differentiate between "didn't run" and "ran
+                               // before command's `end()`
+                               second += 1 + first;
+                             }
+                           });
+
+  scheduler.Schedule(command);
+  EXPECT_EQ(0, first);
+  EXPECT_EQ(0, second);
+  scheduler.Run();
+  EXPECT_EQ(1, first);
+  // if `second == 0`, neither of the lambdas ran.
+  // if `second == 1`, the second lambda ran before the first one
+  EXPECT_EQ(2, second);
+}
+
+// handleInterruptTest() implicitly tests the interrupt=true branch of
+// finallyDo()
+TEST_F(CommandDecoratorTest, HandleInterrupt) {
+  CommandScheduler scheduler = GetScheduler();
+  int first = 0;
+  int second = 0;
+  CommandPtr command = FunctionalCommand([] {}, [] {},
+                                         [&first](bool interrupted) {
+                                           if (interrupted) {
+                                             first++;
+                                           }
+                                         },
+                                         [] { return false; })
+                           .HandleInterrupt([&first, &second] {
+                             // to differentiate between "didn't run" and "ran
+                             // before command's `end()`
+                             second += 1 + first;
+                           });
+
+  scheduler.Schedule(command);
+  scheduler.Run();
+  EXPECT_EQ(0, first);
+  EXPECT_EQ(0, second);
+
+  scheduler.Cancel(command);
+  // if `second == 0`, neither of the lambdas ran.
+  // if `second == 1`, the second lambda ran before the first one
+  EXPECT_EQ(2, second);
+}