Add end effector tests

Signed-off-by: Aryan Khanna <aryankhanna0312@gmail.com>
Change-Id: I13d2da4618eba045ed423c175e02ecef758ed9b6
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