1 files changed, 58 insertions, 0 deletions
diff --git a/cros/update_attempter_unittest.cc b/cros/update_attempter_unittest.cc
index 41da85dd..6d0a1860 100644
--- a/cros/update_attempter_unittest.cc
+++ b/cros/update_attempter_unittest.cc
@@ -22,6 +22,7 @@
 #include <limits>
 #include <map>
 #include <memory>
+#include <optional>
 #include <string>
 #include <unordered_set>
 
@@ -2597,6 +2598,9 @@ TEST_F(UpdateAttempterTest, InvalidateLastPowerwashUpdate) {
   bool save_rollback_data = true;
   FakeSystemState::Get()->fake_hardware()->SchedulePowerwash(
       save_rollback_data);
+  FakeSystemState::Get()
+      ->fake_hardware()
+      ->SetIsPowerwashScheduledByUpdateEngine(true);
   EXPECT_TRUE(FakeSystemState::Get()->fake_hardware()->IsPowerwashScheduled());
   FakeSystemState::Get()->fake_clock()->SetBootTime(Time::FromTimeT(42));
   attempter_.WriteUpdateCompletedMarker();
@@ -2616,6 +2620,60 @@ TEST_F(UpdateAttempterTest, InvalidateLastPowerwashUpdate) {
   EXPECT_TRUE(FakeSystemState::Get()->fake_hardware()->IsFWTryNextSlotReset());
 }
 
+TEST_F(UpdateAttempterTest, InvalidateLastUpdateNoPowerwashFile) {
+  // Mock a previous update.
+  bool save_rollback_data = true;
+  FakeSystemState::Get()->fake_hardware()->SchedulePowerwash(
+      save_rollback_data);
+  EXPECT_TRUE(FakeSystemState::Get()->fake_hardware()->IsPowerwashScheduled());
+  FakeSystemState::Get()
+      ->fake_hardware()
+      ->SetIsPowerwashScheduledByUpdateEngine(std::nullopt);
+  FakeSystemState::Get()->fake_clock()->SetBootTime(Time::FromTimeT(42));
+  attempter_.WriteUpdateCompletedMarker();
+  EXPECT_TRUE(attempter_.GetBootTimeAtUpdate(nullptr));
+  attempter_.status_ = UpdateStatus::UPDATED_NEED_REBOOT;
+  EXPECT_CALL(*FakeSystemState::Get()->mock_payload_state(),
+              SetRollbackHappened(false))
+      .Times(1);
+
+  // Invalidate an update.
+  MockAction action;
+  attempter_.ActionCompleted(
+      nullptr, &action, ErrorCode::kInvalidateLastUpdate);
+
+  EXPECT_FALSE(attempter_.GetBootTimeAtUpdate(nullptr));
+  EXPECT_TRUE(FakeSystemState::Get()->fake_hardware()->IsPowerwashScheduled());
+  EXPECT_TRUE(FakeSystemState::Get()->fake_hardware()->IsFWTryNextSlotReset());
+}
+
+TEST_F(UpdateAttempterTest, InvalidateLastUpdateExternalPowerwash) {
+  // Mock a previous update.
+  bool save_rollback_data = true;
+  FakeSystemState::Get()->fake_hardware()->SchedulePowerwash(
+      save_rollback_data);
+  EXPECT_TRUE(FakeSystemState::Get()->fake_hardware()->IsPowerwashScheduled());
+  FakeSystemState::Get()
+      ->fake_hardware()
+      ->SetIsPowerwashScheduledByUpdateEngine(false);
+  FakeSystemState::Get()->fake_clock()->SetBootTime(Time::FromTimeT(42));
+  attempter_.WriteUpdateCompletedMarker();
+  EXPECT_TRUE(attempter_.GetBootTimeAtUpdate(nullptr));
+  attempter_.status_ = UpdateStatus::UPDATED_NEED_REBOOT;
+  EXPECT_CALL(*FakeSystemState::Get()->mock_payload_state(),
+              SetRollbackHappened(false))
+      .Times(1);
+
+  // Invalidate an update.
+  MockAction action;
+  attempter_.ActionCompleted(
+      nullptr, &action, ErrorCode::kInvalidateLastUpdate);
+
+  EXPECT_FALSE(attempter_.GetBootTimeAtUpdate(nullptr));
+  EXPECT_TRUE(FakeSystemState::Get()->fake_hardware()->IsPowerwashScheduled());
+  EXPECT_TRUE(FakeSystemState::Get()->fake_hardware()->IsFWTryNextSlotReset());
+}
+
 TEST_F(UpdateAttempterTest, CalculateDlcParamsRemoveStaleMetadata) {
   string dlc_id = "dlc0";
   auto active_key =