Add More Checks for SuperstructureState in Superstructure Lib Test

Signed-off-by: Niko Sohmers <nikolai@sohmers.com>
Change-Id: Iedcec7c817fc8f87d42607115efa7c355413083d
diff --git a/y2024/control_loops/superstructure/superstructure.cc b/y2024/control_loops/superstructure/superstructure.cc
index ce5b6a9..56980d5 100644
--- a/y2024/control_loops/superstructure/superstructure.cc
+++ b/y2024/control_loops/superstructure/superstructure.cc
@@ -99,6 +99,7 @@
     if (state_ != SuperstructureState::READY &&
         state_ != SuperstructureState::FIRING) {
       state_ = SuperstructureState::READY;
+      catapult_requested_ = true;
     }
   }
 
@@ -137,7 +138,7 @@
 
   // Checks if the extend is close enough to the retracted position to be
   // considered ready to accept note from the transfer rollers.
-  const bool extend_ready_for_transfer = PositionNear(
+  const bool extend_at_retracted = PositionNear(
       extend_.position(), extend_set_points->retracted(), kExtendThreshold);
 
   // If true, the turret should be moved to the position to avoid collision with
@@ -171,7 +172,7 @@
     case SuperstructureState::IDLE:
       if (unsafe_goal != nullptr &&
           unsafe_goal->intake_goal() == IntakeGoal::INTAKE &&
-          extend_ready_for_transfer) {
+          extend_at_retracted) {
         state_ = SuperstructureState::INTAKING;
       }
       extend_goal = ExtendStatus::RETRACTED;
@@ -282,7 +283,7 @@
             case NoteGoal::NONE:
               extend_goal = ExtendStatus::RETRACTED;
               move_turret_to_standby = true;
-              if (extend_ready_for_transfer) {
+              if (extend_at_retracted) {
                 state_ = SuperstructureState::LOADED;
               }
               break;
diff --git a/y2024/control_loops/superstructure/superstructure_lib_test.cc b/y2024/control_loops/superstructure/superstructure_lib_test.cc
index 1121a5a..ec29ea2 100644
--- a/y2024/control_loops/superstructure/superstructure_lib_test.cc
+++ b/y2024/control_loops/superstructure/superstructure_lib_test.cc
@@ -630,6 +630,10 @@
   RunFor(chrono::seconds(10));
 
   VerifyNearGoal();
+
+  EXPECT_EQ(superstructure_status_fetcher_->state(), SuperstructureState::IDLE);
+  EXPECT_EQ(superstructure_status_fetcher_->shooter()->catapult_state(),
+            CatapultState::READY);
 }
 
 // Tests that loops can reach a goal.
@@ -703,6 +707,9 @@
   RunFor(chrono::seconds(15));
 
   VerifyNearGoal();
+
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::LOADED);
 }
 
 // Makes sure that the voltage on a motor is properly pulled back after
@@ -753,6 +760,9 @@
   RunFor(chrono::seconds(20));
   VerifyNearGoal();
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::READY);
+
   // Try a low acceleration move with a high max velocity and verify the
   // acceleration is capped like expected.
   {
@@ -797,6 +807,9 @@
 
   RunFor(chrono::seconds(10));
   VerifyNearGoal();
+
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::LOADED);
 }
 
 // Tests that the loop zeroes when run for a while without a goal.
@@ -896,6 +909,8 @@
 
   VerifyNearGoal();
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(), SuperstructureState::IDLE);
+
   WaitUntilZeroed();
 
   {
@@ -933,20 +948,8 @@
 
   VerifyNearGoal();
 
-  {
-    auto builder = superstructure_goal_sender_.MakeBuilder();
-
-    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
-
-    goal_builder.add_intake_goal(IntakeGoal::INTAKE);
-    goal_builder.add_note_goal(NoteGoal::NONE);
-
-    ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
-  }
-
-  RunFor(chrono::seconds(5));
-
-  VerifyNearGoal();
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::INTAKING);
 
   EXPECT_EQ(superstructure_output_fetcher_->transfer_roller_voltage(),
             simulated_robot_constants_->common()
@@ -959,6 +962,7 @@
     Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
 
     goal_builder.add_intake_goal(IntakeGoal::INTAKE);
+    goal_builder.add_note_goal(NoteGoal::NONE);
 
     ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
   }
@@ -969,6 +973,9 @@
 
   VerifyNearGoal();
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::LOADED);
+
   EXPECT_EQ(superstructure_output_fetcher_->transfer_roller_voltage(), 0.0);
 }
 
@@ -1012,6 +1019,8 @@
 
   VerifyNearGoal();
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(), SuperstructureState::IDLE);
+
   EXPECT_NEAR(superstructure_status_fetcher_->shooter()->turret()->position(),
               simulated_robot_constants_->common()->turret_loading_position(),
               0.01);
@@ -1047,6 +1056,9 @@
 
   VerifyNearGoal();
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::INTAKING);
+
   EXPECT_NEAR(superstructure_status_fetcher_->shooter()->turret()->position(),
               simulated_robot_constants_->common()->turret_loading_position(),
               0.01);
@@ -1084,11 +1096,15 @@
 
   VerifyNearGoal();
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::INTAKING);
+
   EXPECT_EQ(superstructure_status_fetcher_->shooter()->catapult_state(),
             CatapultState::READY);
 
   EXPECT_EQ(superstructure_status_fetcher_->extend_status(),
             ExtendStatus::RETRACTED);
+
   EXPECT_EQ(superstructure_status_fetcher_->extend_roller(),
             ExtendRollerStatus::TRANSFERING_TO_EXTEND);
 
@@ -1101,6 +1117,9 @@
 
   VerifyNearGoal();
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::LOADED);
+
   EXPECT_EQ(superstructure_status_fetcher_->shooter()->catapult_state(),
             CatapultState::READY);
 
@@ -1142,6 +1161,9 @@
 
   VerifyNearGoal();
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::LOADING_CATAPULT);
+
   EXPECT_EQ(superstructure_status_fetcher_->shooter()->catapult_state(),
             CatapultState::READY);
 
@@ -1158,6 +1180,9 @@
 
   ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::READY);
+
   EXPECT_EQ(superstructure_status_fetcher_->shooter()->catapult_state(),
             CatapultState::LOADED);
 
@@ -1173,9 +1198,6 @@
               simulated_robot_constants_->common()->altitude_loading_position(),
               0.01);
 
-  EXPECT_EQ(superstructure_status_fetcher_->shooter()->catapult_state(),
-            CatapultState::LOADED);
-
   // Fire.  Start by triggering a motion and then firing all in 1 go.
   {
     auto builder = superstructure_goal_sender_.MakeBuilder();
@@ -1219,22 +1241,13 @@
 
   RunFor(chrono::milliseconds(10));
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::FIRING);
+
   // Make sure it stays at firing for a bit.
   EXPECT_EQ(superstructure_status_fetcher_->shooter()->catapult_state(),
             CatapultState::FIRING);
 
-  // Wheel should spin free again.
-  superstructure_plant_.set_catapult_beambreak(false);
-
-  RunFor(chrono::seconds(5));
-
-  // And we should be back to ready.
-  CHECK(superstructure_status_fetcher_.Fetch());
-  EXPECT_EQ(superstructure_status_fetcher_->shooter()->catapult_state(),
-            CatapultState::READY);
-
-  VerifyNearGoal();
-
   EXPECT_NEAR(superstructure_status_fetcher_->shooter()->turret()->position(),
               kTurretGoal, 0.001);
 
@@ -1312,6 +1325,9 @@
   superstructure_status_fetcher_.Fetch();
   ASSERT_TRUE(superstructure_status_fetcher_.get() != nullptr);
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::READY);
+
   EXPECT_EQ(superstructure_status_fetcher_->shooter()->catapult_state(),
             CatapultState::LOADED);
 }
@@ -1381,6 +1397,9 @@
 
   VerifyNearGoal();
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::READY);
+
   EXPECT_NEAR(
       -M_PI_2,
       superstructure_status_fetcher_->shooter()->aimer()->turret_position() -
@@ -1424,6 +1443,9 @@
 
   VerifyNearGoal();
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::READY);
+
   EXPECT_NEAR(
       M_PI_2,
       superstructure_status_fetcher_->shooter()->aimer()->turret_position() +
@@ -1544,6 +1566,9 @@
   ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
   ASSERT_TRUE(superstructure_output_fetcher_.Fetch());
 
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::FIRING);
+
   EXPECT_EQ(superstructure_status_fetcher_->extend_roller(),
             ExtendRollerStatus::SCORING_IN_AMP);
   EXPECT_EQ(superstructure_output_fetcher_->extend_roller_voltage(), 12.0);