[command] Reorder Scheduler operations (#4261)

Fixes several cases where calling scheduler operations from a command callback could result in NPEs or other issues.

Co-authored-by: Tyler Veness <calcmogul@gmail.com>

wpilibNewCommands/src/test/native/cpp/frc2/command/CommandTestBase.h

@@ -18,6 +18,7 @@
 #include "make_vector.h"
 
 namespace frc2 {
+
 class CommandTestBase : public ::testing::Test {
  public:
   CommandTestBase();
@@ -93,4 +94,91 @@ class CommandTestBase : public ::testing::Test {
 
   void SetDSEnabled(bool enabled);
 };
+
+template <typename T>
+class CommandTestBaseWithParam : public ::testing::TestWithParam<T> {
+ public:
+  CommandTestBaseWithParam() {
+    auto& scheduler = CommandScheduler::GetInstance();
+    scheduler.CancelAll();
+    scheduler.Enable();
+    scheduler.GetActiveButtonLoop()->Clear();
+  }
+
+  class TestSubsystem : public SubsystemBase {};
+
+ protected:
+  class MockCommand : public Command {
+   public:
+    MOCK_CONST_METHOD0(GetRequirements, wpi::SmallSet<Subsystem*, 4>());
+    MOCK_METHOD0(IsFinished, bool());
+    MOCK_CONST_METHOD0(RunsWhenDisabled, bool());
+    MOCK_METHOD0(Initialize, void());
+    MOCK_METHOD0(Execute, void());
+    MOCK_METHOD1(End, void(bool interrupted));
+
+    MockCommand() {
+      m_requirements = {};
+      EXPECT_CALL(*this, GetRequirements())
+          .WillRepeatedly(::testing::Return(m_requirements));
+      EXPECT_CALL(*this, IsFinished()).WillRepeatedly(::testing::Return(false));
+      EXPECT_CALL(*this, RunsWhenDisabled())
+          .WillRepeatedly(::testing::Return(true));
+    }
+
+    MockCommand(std::initializer_list<Subsystem*> requirements,
+                bool finished = false, bool runWhenDisabled = true) {
+      m_requirements.insert(requirements.begin(), requirements.end());
+      EXPECT_CALL(*this, GetRequirements())
+          .WillRepeatedly(::testing::Return(m_requirements));
+      EXPECT_CALL(*this, IsFinished())
+          .WillRepeatedly(::testing::Return(finished));
+      EXPECT_CALL(*this, RunsWhenDisabled())
+          .WillRepeatedly(::testing::Return(runWhenDisabled));
+    }
+
+    MockCommand(MockCommand&& other) {
+      EXPECT_CALL(*this, IsFinished())
+          .WillRepeatedly(::testing::Return(other.IsFinished()));
+      EXPECT_CALL(*this, RunsWhenDisabled())
+          .WillRepeatedly(::testing::Return(other.RunsWhenDisabled()));
+      std::swap(m_requirements, other.m_requirements);
+      EXPECT_CALL(*this, GetRequirements())
+          .WillRepeatedly(::testing::Return(m_requirements));
+    }
+
+    MockCommand(const MockCommand& other) : Command{other} {}
+
+    void SetFinished(bool finished) {
+      EXPECT_CALL(*this, IsFinished())
+          .WillRepeatedly(::testing::Return(finished));
+    }
+
+    ~MockCommand() {  // NOLINT
+      auto& scheduler = CommandScheduler::GetInstance();
+      scheduler.Cancel(this);
+    }
+
+   protected:
+    std::unique_ptr<Command> TransferOwnership() && {  // NOLINT
+      return std::make_unique<MockCommand>(std::move(*this));
+    }
+
+   private:
+    wpi::SmallSet<Subsystem*, 4> m_requirements;
+  };
+
+  CommandScheduler GetScheduler() { return CommandScheduler(); }
+
+  void SetUp() override { frc::sim::DriverStationSim::SetEnabled(true); }
+
+  void TearDown() override {
+    CommandScheduler::GetInstance().GetActiveButtonLoop()->Clear();
+  }
+
+  void SetDSEnabled(bool enabled) {
+    frc::sim::DriverStationSim::SetEnabled(enabled);
+  }
+};
+
 }  // namespace frc2