Add end effector tests

Signed-off-by: Aryan Khanna <aryankhanna0312@gmail.com>
Change-Id: I13d2da4618eba045ed423c175e02ecef758ed9b6
diff --git a/y2023_bot3/constants.h b/y2023_bot3/constants.h
index 500b836..6cb7d8e 100644
--- a/y2023_bot3/constants.h
+++ b/y2023_bot3/constants.h
@@ -39,6 +39,9 @@
   static constexpr double kDrivetrainSupplyCurrentLimit() { return 35.0; }
   static constexpr double kDrivetrainStatorCurrentLimit() { return 60.0; }
 
+  static constexpr double kRollerSupplyCurrentLimit() { return 30.0; }
+  static constexpr double kRollerStatorCurrentLimit() { return 100.0; }
+
   // timeout to ensure code doesn't get stuck after releasing the "intake"
   // button
   static constexpr std::chrono::milliseconds kExtraIntakingTime() {
diff --git a/y2023_bot3/control_loops/superstructure/end_effector.cc b/y2023_bot3/control_loops/superstructure/end_effector.cc
index a98d06f..178e946 100644
--- a/y2023_bot3/control_loops/superstructure/end_effector.cc
+++ b/y2023_bot3/control_loops/superstructure/end_effector.cc
@@ -18,14 +18,6 @@
     bool beambreak_status, double *roller_voltage, bool preloaded_with_cube) {
   *roller_voltage = 0.0;
 
-  if (roller_goal == RollerGoal::SPIT) {
-    state_ = EndEffectorState::SPITTING;
-  }
-
-  if (roller_goal == RollerGoal::SPIT_MID) {
-    state_ = EndEffectorState::SPITTING_MID;
-  }
-
   // If we started off preloaded, skip to the loaded state.
   // Make sure we weren't already there just in case.
   if (preloaded_with_cube) {
@@ -43,6 +35,14 @@
     }
   }
 
+  if (roller_goal == RollerGoal::SPIT) {
+    state_ = EndEffectorState::SPITTING;
+  }
+
+  if (roller_goal == RollerGoal::SPIT_MID) {
+    state_ = EndEffectorState::SPITTING_MID;
+  }
+
   switch (state_) {
     case EndEffectorState::IDLE:
       // If idle and intake requested, intake
@@ -72,12 +72,30 @@
       break;
     case EndEffectorState::LOADED:
       timer_ = timestamp;
+      if (!preloaded_with_cube && !beambreak_status) {
+        state_ = EndEffectorState::INTAKING;
+        break;
+      }
+
       break;
+
     case EndEffectorState::SPITTING:
       *roller_voltage = kRollerCubeSpitVoltage();
+
+      if (roller_goal == RollerGoal::IDLE) {
+        state_ = EndEffectorState::IDLE;
+      }
+
       break;
+
     case EndEffectorState::SPITTING_MID:
       *roller_voltage = kRollerCubeSpitMidVoltage();
+
+      if (roller_goal == RollerGoal::IDLE) {
+        state_ = EndEffectorState::IDLE;
+      }
+
+      break;
   }
 }
 
diff --git a/y2023_bot3/control_loops/superstructure/superstructure.cc b/y2023_bot3/control_loops/superstructure/superstructure.cc
index e9df534..740f4a8 100644
--- a/y2023_bot3/control_loops/superstructure/superstructure.cc
+++ b/y2023_bot3/control_loops/superstructure/superstructure.cc
@@ -33,25 +33,29 @@
                                   const Position *position,
                                   aos::Sender<Output>::Builder *output,
                                   aos::Sender<Status>::Builder *status) {
-  (void)unsafe_goal;
-  (void)position;
-
   const monotonic_clock::time_point timestamp =
       event_loop()->context().monotonic_event_time;
-  (void)timestamp;
 
   if (WasReset()) {
     AOS_LOG(ERROR, "WPILib reset, restarting\n");
+    end_effector_.Reset();
   }
 
   OutputT output_struct;
 
+  end_effector_.RunIteration(
+      timestamp,
+      unsafe_goal != nullptr ? unsafe_goal->roller_goal() : RollerGoal::IDLE,
+      position->end_effector_cube_beam_break(), &output_struct.roller_voltage,
+      unsafe_goal != nullptr ? unsafe_goal->preloaded_with_cube() : false);
+
   if (output) {
     output->CheckOk(output->Send(Output::Pack(*output->fbb(), &output_struct)));
   }
 
   Status::Builder status_builder = status->MakeBuilder<Status>();
   status_builder.add_zeroed(true);
+  status_builder.add_end_effector_state(end_effector_.state());
 
   (void)status->Send(status_builder.Finish());
 }
diff --git a/y2023_bot3/control_loops/superstructure/superstructure_lib_test.cc b/y2023_bot3/control_loops/superstructure/superstructure_lib_test.cc
index 24c4f2a..7387ca7 100644
--- a/y2023_bot3/control_loops/superstructure/superstructure_lib_test.cc
+++ b/y2023_bot3/control_loops/superstructure/superstructure_lib_test.cc
@@ -66,14 +66,22 @@
         superstructure_position_sender_.MakeBuilder();
 
     Position::Builder position_builder = builder.MakeBuilder<Position>();
+    position_builder.add_end_effector_cube_beam_break(
+        end_effector_cube_beam_break_);
     CHECK_EQ(builder.Send(position_builder.Finish()),
              aos::RawSender::Error::kOk);
   }
 
+  void set_end_effector_cube_beam_break(bool triggered) {
+    end_effector_cube_beam_break_ = triggered;
+  }
+
  private:
   ::aos::EventLoop *event_loop_;
   ::aos::PhasedLoopHandler *phased_loop_handle_ = nullptr;
 
+  bool end_effector_cube_beam_break_ = false;
+
   ::aos::Sender<Position> superstructure_position_sender_;
   ::aos::Fetcher<Status> superstructure_status_fetcher_;
   ::aos::Fetcher<Output> superstructure_output_fetcher_;
@@ -242,6 +250,211 @@
   CheckIfZeroed();
 }
 
+TEST_F(SuperstructureTest, EndEffectorGoal) {
+  SetEnabled(true);
+  WaitUntilZeroed();
+
+  double spit_voltage = EndEffector::kRollerCubeSpitVoltage();
+  double suck_voltage = EndEffector::kRollerCubeSuckVoltage();
+
+  RollerGoal roller_goal = RollerGoal::INTAKE_CUBE;
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_roller_goal(roller_goal);
+
+    builder.CheckOk(builder.Send(goal_builder.Finish()));
+  }
+  superstructure_plant_.set_end_effector_cube_beam_break(false);
+
+  // This makes sure that we intake as normal when
+  // requesting intake.
+  RunFor(constants::Values::kExtraIntakingTime());
+
+  ASSERT_TRUE(superstructure_output_fetcher_.Fetch());
+  ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
+
+  EXPECT_EQ(superstructure_output_fetcher_->roller_voltage(), suck_voltage);
+  EXPECT_EQ(superstructure_status_fetcher_->end_effector_state(),
+            EndEffectorState::INTAKING);
+
+  superstructure_plant_.set_end_effector_cube_beam_break(true);
+
+  // Checking that after the beambreak is set once intaking that the
+  // state changes to LOADED.
+  RunFor(dt());
+
+  ASSERT_TRUE(superstructure_output_fetcher_.Fetch());
+  ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
+
+  EXPECT_EQ(superstructure_output_fetcher_->roller_voltage(), 0.0);
+  EXPECT_EQ(superstructure_status_fetcher_->end_effector_state(),
+            EndEffectorState::LOADED);
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_roller_goal(RollerGoal::IDLE);
+
+    builder.CheckOk(builder.Send(goal_builder.Finish()));
+  }
+  superstructure_plant_.set_end_effector_cube_beam_break(false);
+
+  //  Checking that it's going back to intaking because we lost the
+  //  beambreak sensor.
+  RunFor(dt() * 2);
+
+  ASSERT_TRUE(superstructure_output_fetcher_.Fetch());
+  ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
+
+  EXPECT_EQ(superstructure_output_fetcher_->roller_voltage(), suck_voltage);
+  EXPECT_EQ(superstructure_status_fetcher_->end_effector_state(),
+            EndEffectorState::INTAKING);
+
+  // Checking that we go back to idle after beambreak is lost and we
+  // set our goal to idle.
+  RunFor(dt() * 2 + constants::Values::kExtraIntakingTime());
+  ASSERT_TRUE(superstructure_output_fetcher_.Fetch());
+  ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
+
+  EXPECT_EQ(superstructure_output_fetcher_->roller_voltage(), 0.0);
+  EXPECT_EQ(superstructure_status_fetcher_->end_effector_state(),
+            EndEffectorState::IDLE);
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_roller_goal(roller_goal);
+
+    builder.CheckOk(builder.Send(goal_builder.Finish()));
+  }
+
+  // Going through intake -> loaded -> spitting
+  // Making sure that it's intaking normally.
+  RunFor(constants::Values::kExtraIntakingTime());
+
+  ASSERT_TRUE(superstructure_output_fetcher_.Fetch());
+  ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
+
+  EXPECT_EQ(superstructure_output_fetcher_->roller_voltage(), suck_voltage);
+  EXPECT_EQ(superstructure_status_fetcher_->end_effector_state(),
+            EndEffectorState::INTAKING);
+
+  superstructure_plant_.set_end_effector_cube_beam_break(true);
+
+  // Checking that it's loaded once beambreak is sensing something.
+  RunFor(dt());
+
+  ASSERT_TRUE(superstructure_output_fetcher_.Fetch());
+  ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
+
+  EXPECT_EQ(superstructure_output_fetcher_->roller_voltage(), 0.0);
+  EXPECT_EQ(superstructure_status_fetcher_->end_effector_state(),
+            EndEffectorState::LOADED);
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_roller_goal(RollerGoal::SPIT);
+
+    builder.CheckOk(builder.Send(goal_builder.Finish()));
+  }
+  superstructure_plant_.set_end_effector_cube_beam_break(true);
+  // Checking that it stays spitting until 2 seconds after the
+  // beambreak is lost.
+  RunFor(dt() * 10);
+
+  ASSERT_TRUE(superstructure_output_fetcher_.Fetch());
+  ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
+
+  EXPECT_EQ(superstructure_output_fetcher_->roller_voltage(), spit_voltage);
+  EXPECT_EQ(superstructure_status_fetcher_->end_effector_state(),
+            EndEffectorState::SPITTING);
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+
+    goal_builder.add_roller_goal(RollerGoal::IDLE);
+
+    builder.CheckOk(builder.Send(goal_builder.Finish()));
+  }
+
+  // Checking that it goes to idle after it's given time to stop spitting.
+  RunFor(dt() * 3);
+
+  ASSERT_TRUE(superstructure_output_fetcher_.Fetch());
+  ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
+
+  EXPECT_EQ(superstructure_output_fetcher_->roller_voltage(), 0.0);
+  EXPECT_EQ(superstructure_status_fetcher_->end_effector_state(),
+            EndEffectorState::IDLE);
+}
+
+// Test that we are able to signal that the cube was preloaded
+TEST_F(SuperstructureTest, Preloaded) {
+  SetEnabled(true);
+  WaitUntilZeroed();
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_preloaded_with_cube(true);
+    ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
+  }
+
+  RunFor(dt());
+
+  ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
+  EXPECT_EQ(superstructure_status_fetcher_->end_effector_state(),
+            EndEffectorState::LOADED);
+}
+
+// Tests that the end effector does nothing when the goal is to remain
+// still.
+TEST_F(SuperstructureTest, DoesNothing) {
+  SetEnabled(true);
+  WaitUntilZeroed();
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+
+    goal_builder.add_roller_goal(RollerGoal::IDLE);
+
+    ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
+  }
+  RunFor(chrono::seconds(10));
+  VerifyNearGoal();
+
+  EXPECT_TRUE(superstructure_output_fetcher_.Fetch());
+}
+// Tests that loops can reach a goal.
+TEST_F(SuperstructureTest, ReachesGoal) {
+  SetEnabled(true);
+  WaitUntilZeroed();
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+
+    goal_builder.add_roller_goal(RollerGoal::IDLE);
+
+    ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
+  }
+  // Give it a lot of time to get there.
+  RunFor(chrono::seconds(15));
+
+  VerifyNearGoal();
+}
+
 }  // namespace testing
 }  // namespace superstructure
 }  // namespace control_loops
diff --git a/y2023_bot3/wpilib_interface.cc b/y2023_bot3/wpilib_interface.cc
index 2b5a9f4..09f87db 100644
--- a/y2023_bot3/wpilib_interface.cc
+++ b/y2023_bot3/wpilib_interface.cc
@@ -177,6 +177,34 @@
     PrintConfigs();
   }
 
+  void WriteRollerConfigs() {
+    ctre::phoenix6::configs::CurrentLimitsConfigs current_limits;
+    current_limits.StatorCurrentLimit =
+        constants::Values::kRollerStatorCurrentLimit();
+    current_limits.StatorCurrentLimitEnable = true;
+    current_limits.SupplyCurrentLimit =
+        constants::Values::kRollerSupplyCurrentLimit();
+    current_limits.SupplyCurrentLimitEnable = true;
+
+    ctre::phoenix6::configs::MotorOutputConfigs output_configs;
+    output_configs.NeutralMode =
+        ctre::phoenix6::signals::NeutralModeValue::Brake;
+    output_configs.DutyCycleNeutralDeadband = 0;
+
+    ctre::phoenix6::configs::TalonFXConfiguration configuration;
+    configuration.CurrentLimits = current_limits;
+    configuration.MotorOutput = output_configs;
+
+    ctre::phoenix::StatusCode status =
+        talon_.GetConfigurator().Apply(configuration);
+    if (!status.IsOK()) {
+      AOS_LOG(ERROR, "Failed to set falcon configuration: %s: %s",
+              status.GetName(), status.GetDescription());
+    }
+
+    PrintConfigs();
+  }
+
   ctre::phoenix6::hardware::TalonFX *talon() { return &talon_; }
 
   flatbuffers::Offset<frc971::control_loops::CANFalcon> WritePosition(
@@ -243,7 +271,8 @@
         can_position_sender_(
             event_loop
                 ->MakeSender<frc971::control_loops::drivetrain::CANPosition>(
-                    "/drivetrain")) {
+                    "/drivetrain")),
+        roller_falcon_data_(std::nullopt) {
     event_loop->SetRuntimeRealtimePriority(40);
     event_loop->SetRuntimeAffinity(aos::MakeCpusetFromCpus({1}));
     timer_handler_ = event_loop->AddTimer([this]() { Loop(); });
@@ -264,7 +293,7 @@
     right_back_ = std::move(right_back);
     left_front_ = std::move(left_front);
     left_back_ = std::move(left_back);
-    roller_falcon = std::move(roller_falcon);
+    roller_falcon_ = std::move(roller_falcon);
   }
 
   std::optional<frc971::control_loops::CANFalconT> roller_falcon_data() {
@@ -284,7 +313,8 @@
 
     auto builder = can_position_sender_.MakeBuilder();
 
-    for (auto falcon : {right_front_, right_back_, left_front_, left_back_}) {
+    for (auto falcon :
+         {right_front_, right_back_, left_front_, left_back_, roller_falcon_}) {
       falcon->RefreshNontimesyncedSignals();
     }
 
@@ -311,6 +341,21 @@
     can_position_builder.add_status(static_cast<int>(status));
 
     builder.CheckOk(builder.Send(can_position_builder.Finish()));
+
+    {
+      std::unique_lock<aos::stl_mutex> lock(roller_mutex_);
+      frc971::control_loops::CANFalconT roller_falcon_data;
+      roller_falcon_data.id = roller_falcon_->device_id();
+      roller_falcon_data.supply_current = roller_falcon_->supply_current();
+      roller_falcon_data.torque_current = -roller_falcon_->torque_current();
+      roller_falcon_data.supply_voltage = roller_falcon_->supply_voltage();
+      roller_falcon_data.device_temp = roller_falcon_->device_temp();
+      roller_falcon_data.position = -roller_falcon_->position();
+      roller_falcon_data.duty_cycle = roller_falcon_->duty_cycle();
+      roller_falcon_data_ =
+          std::make_optional<frc971::control_loops::CANFalconT>(
+              roller_falcon_data);
+    }
   }
 
   aos::EventLoop *event_loop_;
@@ -320,7 +365,7 @@
       can_position_sender_;
 
   std::shared_ptr<Falcon> right_front_, right_back_, left_front_, left_back_,
-      roller_falcon;
+      roller_falcon_;
 
   std::optional<frc971::control_loops::CANFalconT> roller_falcon_data_;
 
@@ -385,8 +430,7 @@
       superstructure::Position::Builder position_builder =
           builder.MakeBuilder<superstructure::Position>();
       position_builder.add_end_effector_cube_beam_break(
-          end_effector_cube_beam_break_->Get());
-
+          !end_effector_cube_beam_break_->Get());
       if (!roller_falcon_offset.IsNull()) {
         position_builder.add_roller_falcon(roller_falcon_offset);
       }
@@ -468,6 +512,11 @@
     superstructure_reading_ = superstructure_reading;
   }
 
+  void set_end_effector_cube_beam_break(
+      ::std::unique_ptr<frc::DigitalInput> sensor) {
+    end_effector_cube_beam_break_ = ::std::move(sensor);
+  }
+
  private:
   std::shared_ptr<const Values> values_;
 
@@ -536,26 +585,48 @@
   };
 
   void HandleCANConfiguration(const frc971::CANConfiguration &configuration) {
+    roller_falcon_->PrintConfigs();
     if (configuration.reapply()) {
       WriteConfigs();
     }
   }
 
- private:
-  void WriteConfigs() {}
+  void set_roller_falcon(std::shared_ptr<Falcon> roller_falcon) {
+    roller_falcon_ = std::move(roller_falcon);
+  }
 
-  void Write(const superstructure::Output &output) override { (void)output; }
+ private:
+  void WriteConfigs() { roller_falcon_->WriteRollerConfigs(); }
+
+  void Write(const superstructure::Output &output) override {
+    ctre::phoenix6::controls::DutyCycleOut roller_control(
+        SafeSpeed(-output.roller_voltage()));
+    roller_control.UpdateFreqHz = 0_Hz;
+    roller_control.EnableFOC = true;
+
+    ctre::phoenix::StatusCode status =
+        roller_falcon_->talon()->SetControl(roller_control);
+
+    if (!status.IsOK()) {
+      AOS_LOG(ERROR, "Failed to write control to falcon: %s: %s",
+              status.GetName(), status.GetDescription());
+    }
+  }
 
   void Stop() override {
     AOS_LOG(WARNING, "Superstructure CAN output too old.\n");
     ctre::phoenix6::controls::DutyCycleOut stop_command(0.0);
     stop_command.UpdateFreqHz = 0_Hz;
     stop_command.EnableFOC = true;
+
+    roller_falcon_->talon()->SetControl(stop_command);
   }
 
   double SafeSpeed(double voltage) {
     return (::aos::Clip(voltage, -kMaxBringupPower, kMaxBringupPower) / 12.0);
   }
+
+  std::shared_ptr<Falcon> roller_falcon_;
 };
 
 class DrivetrainWriter : public ::frc971::wpilib::LoopOutputHandler<
@@ -694,13 +765,13 @@
     std::shared_ptr<Falcon> right_front =
         std::make_shared<Falcon>(1, "Drivetrain Bus", &signals_registry);
     std::shared_ptr<Falcon> right_back =
-        std::make_shared<Falcon>(2, "Drivetrain Bus", &signals_registry);
+        std::make_shared<Falcon>(0, "Drivetrain Bus", &signals_registry);
     std::shared_ptr<Falcon> left_front =
-        std::make_shared<Falcon>(4, "Drivetrain Bus", &signals_registry);
+        std::make_shared<Falcon>(2, "Drivetrain Bus", &signals_registry);
     std::shared_ptr<Falcon> left_back =
+        std::make_shared<Falcon>(3, "Drivetrain Bus", &signals_registry);
+    std::shared_ptr<Falcon> roller =
         std::make_shared<Falcon>(5, "Drivetrain Bus", &signals_registry);
-    std::shared_ptr<Falcon> roller_falcon =
-        std::make_shared<Falcon>(13, "Drivetrain Bus", &signals_registry);
 
     // Thread 3.
     ::aos::ShmEventLoop can_sensor_reader_event_loop(&config.message());
@@ -709,7 +780,7 @@
                                       std::move(signals_registry));
 
     can_sensor_reader.set_falcons(right_front, right_back, left_front,
-                                  left_back, roller_falcon);
+                                  left_back, roller);
 
     AddLoop(&can_sensor_reader_event_loop);
 
@@ -722,6 +793,8 @@
     sensor_reader.set_drivetrain_right_encoder(make_encoder(0));
     sensor_reader.set_superstructure_reading(superstructure_reading);
     sensor_reader.set_yaw_rate_input(make_unique<frc::DigitalInput>(0));
+    sensor_reader.set_end_effector_cube_beam_break(
+        make_unique<frc::DigitalInput>(22));
 
     AddLoop(&sensor_reader_event_loop);
 
@@ -767,12 +840,18 @@
     AddLoop(&output_event_loop);
 
     // Thread 7
-    // Set up led_indicator.
-    ::aos::ShmEventLoop led_indicator_event_loop(&config.message());
-    led_indicator_event_loop.set_name("LedIndicator");
-    control_loops::superstructure::LedIndicator led_indicator(
-        &led_indicator_event_loop);
-    AddLoop(&led_indicator_event_loop);
+    // Setup superstructure CAN output.
+    SuperstructureCANWriter superstructure_can_writer(&can_output_event_loop);
+    superstructure_can_writer.set_roller_falcon(roller);
+
+    can_output_event_loop.MakeWatcher(
+        "/roborio", [&drivetrain_writer, &superstructure_can_writer](
+                        const frc971::CANConfiguration &configuration) {
+          drivetrain_writer.HandleCANConfiguration(configuration);
+          superstructure_can_writer.HandleCANConfiguration(configuration);
+        });
+
+    AddLoop(&can_output_event_loop);
 
     RunLoops();
   }