Add unittest for resume update

Test: treehugger
Bug: 173392916
Change-Id: I4b320685af93b20fc1f4abcdd6a416f1c97731ab
diff --git a/payload_consumer/delta_performer.cc b/payload_consumer/delta_performer.cc
index c3a321e..be39542 100644
--- a/payload_consumer/delta_performer.cc
+++ b/payload_consumer/delta_performer.cc
@@ -1443,7 +1443,14 @@
       TEST_AND_RETURN_FALSE(
           prefs_->SetInt64(kPrefsUpdateStateNextDataLength, 0));
     }
-    partition_writer_->CheckpointUpdateProgress(GetPartitionOperationNum());
+    if (partition_writer_) {
+      partition_writer_->CheckpointUpdateProgress(GetPartitionOperationNum());
+    } else {
+      CHECK_EQ(next_operation_num_, num_total_operations_)
+          << "Partition writer is null, we are expected to finish all "
+             "operations: "
+          << next_operation_num_ << "/" << num_total_operations_;
+    }
   }
   TEST_AND_RETURN_FALSE(
       prefs_->SetInt64(kPrefsUpdateStateNextOperation, next_operation_num_));
diff --git a/payload_consumer/delta_performer.h b/payload_consumer/delta_performer.h
index 7df2bd3..67ca6ce 100644
--- a/payload_consumer/delta_performer.h
+++ b/payload_consumer/delta_performer.h
@@ -170,6 +170,12 @@
   // Return true if header parsing is finished and no errors occurred.
   bool IsHeaderParsed() const;
 
+  // Checkpoints the update progress into persistent storage to allow this
+  // update attempt to be resumed after reboot.
+  // If |force| is false, checkpoint may be throttled.
+  // Exposed for testing purposes.
+  bool CheckpointUpdateProgress(bool force);
+
   // Compare |calculated_hash| with source hash in |operation|, return false and
   // dump hash and set |error| if don't match.
   // |source_fd| is the file descriptor of the source partition.
@@ -272,11 +278,6 @@
   // accordingly.
   void DiscardBuffer(bool do_advance_offset, size_t signed_hash_buffer_size);
 
-  // Checkpoints the update progress into persistent storage to allow this
-  // update attempt to be resumed after reboot.
-  // If |force| is false, checkpoint may be throttled.
-  bool CheckpointUpdateProgress(bool force);
-
   // Primes the required update state. Returns true if the update state was
   // successfully initialized to a saved resume state or if the update is a new
   // update. Returns false otherwise.
diff --git a/payload_consumer/delta_performer_unittest.cc b/payload_consumer/delta_performer_unittest.cc
index 74f17d2..ba387a5 100644
--- a/payload_consumer/delta_performer_unittest.cc
+++ b/payload_consumer/delta_performer_unittest.cc
@@ -59,7 +59,6 @@
 using std::vector;
 using test_utils::GetBuildArtifactsPath;
 using test_utils::kRandomString;
-using test_utils::System;
 using testing::_;
 
 extern const char* kUnittestPrivateKeyPath;
@@ -418,8 +417,6 @@
     EXPECT_EQ(payload_.metadata_size, performer_.metadata_size_);
   }
 
-
-
   FakePrefs prefs_;
   InstallPlan install_plan_;
   InstallPlan::Payload payload_;
@@ -1072,4 +1069,30 @@
   EXPECT_EQ(kMaxSupportedMajorPayloadVersion, major_version);
 }
 
+TEST_F(DeltaPerformerTest, FullPayloadCanResumeTest) {
+  payload_.type = InstallPayloadType::kFull;
+  brillo::Blob expected_data =
+      brillo::Blob(std::begin(kRandomString), std::end(kRandomString));
+  expected_data.resize(4096);  // block size
+  vector<AnnotatedOperation> aops;
+  AnnotatedOperation aop;
+  *(aop.op.add_dst_extents()) = ExtentForRange(0, 1);
+  aop.op.set_data_offset(0);
+  aop.op.set_data_length(expected_data.size());
+  aop.op.set_type(InstallOperation::REPLACE);
+  aops.push_back(aop);
+
+  brillo::Blob payload_data = GeneratePayload(expected_data,
+                                              aops,
+                                              false,
+                                              kBrilloMajorPayloadVersion,
+                                              kFullPayloadMinorVersion);
+
+  ASSERT_EQ(expected_data, ApplyPayload(payload_data, "/dev/null", true));
+  performer_.CheckpointUpdateProgress(true);
+  const std::string payload_id = "12345";
+  prefs_.SetString(kPrefsUpdateCheckResponseHash, payload_id);
+  ASSERT_TRUE(DeltaPerformer::CanResumeUpdate(&prefs_, payload_id));
+}
+
 }  // namespace chromeos_update_engine