Make the BattOr agent retry start tracing when we fail to read an ack

tools/battor_agent/battor_agent_unittest.cc

Index: tools/battor_agent/battor_agent_unittest.cc
diff --git a/tools/battor_agent/battor_agent_unittest.cc b/tools/battor_agent/battor_agent_unittest.cc
index 75ec3015930070b1505461e0eacc8a16b44658be..291c109c9ae016f7a3620e3bedbf1678cca3e631 100644
--- a/tools/battor_agent/battor_agent_unittest.cc
+++ b/tools/battor_agent/battor_agent_unittest.cc
@@ -120,7 +120,7 @@ class BattOrAgentTest : public testing::Test, public BattOrAgent::Listener {
   }
 
   void OnGetFirmwareGitHashComplete(const std::string& firmware_git_hash,
-                            BattOrError error) override {
+                                    BattOrError error) override {
     is_command_complete_ = true;
     command_error_ = error;
     firmware_git_hash_ = firmware_git_hash;
@@ -401,7 +401,7 @@ TEST_F(BattOrAgentTest, StartTracingFailsIfInitAckReadFails) {
 
     // Bytes will be sent because INIT will be retried.
     OnBytesSent(true);
-   }
+  }
 
   EXPECT_TRUE(IsCommandComplete());
   EXPECT_EQ(BATTOR_ERROR_TOO_MANY_INIT_RETRIES, GetCommandError());
@@ -455,20 +455,86 @@ TEST_F(BattOrAgentTest, StartTracingFailsIfStartTracingSendFails) {
   EXPECT_EQ(BATTOR_ERROR_SEND_ERROR, GetCommandError());
 }
 
-TEST_F(BattOrAgentTest, StartTracingFailsIfStartTracingAckReadFails) {
-  RunStartTracingTo(BattOrAgentState::START_TRACING_SENT);
-  OnMessageRead(false, BATTOR_MESSAGE_TYPE_CONTROL_ACK, nullptr);
+TEST_F(BattOrAgentTest, StartTracingSucceedsAfterRetriesIfWrongAckRead) {
+  RunStartTracingTo(BattOrAgentState::CONNECTED);
+
+  for (int i = 0; i < 4; i++) {
+    // Go through the correct init sequence, but give the wrong ack to
+    // START_TRACING.
+    OnBytesSent(true);
+    OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK,
+                  ToCharVector(kInitAck));
+    OnBytesSent(true);
+    OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK,
+                  ToCharVector(kSetGainAck));
+    OnBytesSent(true);
+    OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK,
+                  ToCharVector(kInitAck));
+  }
+
+  // On the last attempt, give the correct ack to START_TRACING.
+  OnBytesSent(true);
+  OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK, ToCharVector(kInitAck));
+  OnBytesSent(true);
+  OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK,
+                ToCharVector(kSetGainAck));
+  OnBytesSent(true);
+  OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK,
+                ToCharVector(kStartTracingAck));
 
   EXPECT_TRUE(IsCommandComplete());
-  EXPECT_EQ(BATTOR_ERROR_RECEIVE_ERROR, GetCommandError());
+  EXPECT_EQ(BATTOR_ERROR_NONE, GetCommandError());
 }
 
-TEST_F(BattOrAgentTest, StartTracingFailsIfStartTracingWrongAckRead) {
-  RunStartTracingTo(BattOrAgentState::START_TRACING_SENT);
+TEST_F(BattOrAgentTest, StartTracingSucceedsAfterRetriesWithReadFailure) {
+  RunStartTracingTo(BattOrAgentState::CONNECTED);
+
+  for (int i = 0; i < 4; i++) {
+    // Go through the correct init sequence, but indicate that we failed to read
+    // the START_TRACING ack.
+    OnBytesSent(true);
+    OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK,
+                  ToCharVector(kInitAck));
+    OnBytesSent(true);
+    OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK,
+                  ToCharVector(kSetGainAck));
+    OnBytesSent(true);
+    OnMessageRead(false, BATTOR_MESSAGE_TYPE_CONTROL_ACK, nullptr);
+  }
+
+  // On the last attempt, give the correct ack to START_TRACING.
+  OnBytesSent(true);
   OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK, ToCharVector(kInitAck));
+  OnBytesSent(true);
+  OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK,
+                ToCharVector(kSetGainAck));
+  OnBytesSent(true);
+  OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK,
+                ToCharVector(kStartTracingAck));
 
   EXPECT_TRUE(IsCommandComplete());
-  EXPECT_EQ(BATTOR_ERROR_UNEXPECTED_MESSAGE, GetCommandError());
+  EXPECT_EQ(BATTOR_ERROR_NONE, GetCommandError());
+}
+
+TEST_F(BattOrAgentTest, StartTracingFailsIfStartTracingWrongAckReadTooMuch) {
+  RunStartTracingTo(BattOrAgentState::CONNECTED);
+
+  for (int i = 0; i < 5; i++) {
+    // Go through the correct init sequence, but give the wrong ack to
+    // START_TRACING.
+    OnBytesSent(true);
+    OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK,
+                  ToCharVector(kInitAck));
+    OnBytesSent(true);
+    OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK,
+                  ToCharVector(kSetGainAck));
+    OnBytesSent(true);
+    OnMessageRead(true, BATTOR_MESSAGE_TYPE_CONTROL_ACK,
+                  ToCharVector(kInitAck));
+  }
+
+  EXPECT_TRUE(IsCommandComplete());
+  EXPECT_EQ(BATTOR_ERROR_TOO_MANY_START_TRACING_RETRIES, GetCommandError());
 }
 
 TEST_F(BattOrAgentTest, StartTracingSucceedsWithOneInitFailure) {
@@ -887,7 +953,6 @@ TEST_F(BattOrAgentTest, RecordClockSyncMarkerPrintsInStopTracingResult) {
       "0.00 0.0 0.0 <MY_MARKER>\n"
       "1.00 0.6 1.2\n",
       GetTrace());
-
 }
 
 TEST_F(BattOrAgentTest, RecordClockSyncMarkerFailsWithoutConnection) {