diff --git a/payload_consumer/delta_performer.cc b/payload_consumer/delta_performer.cc
index be39542..a10815b 100644
--- a/payload_consumer/delta_performer.cc
+++ b/payload_consumer/delta_performer.cc
@@ -240,13 +240,13 @@
   const InstallPlan::Partition& install_part =
       install_plan_->partitions[num_previous_partitions + current_partition_];
   auto dynamic_control = boot_control_->GetDynamicPartitionControl();
-  partition_writer_ = partition_writer::CreatePartitionWriter(
-      partition,
-      install_part,
-      dynamic_control,
-      block_size_,
-      interactive_,
-      IsDynamicPartition(install_part.name));
+  partition_writer_ =
+      CreatePartitionWriter(partition,
+                            install_part,
+                            dynamic_control,
+                            block_size_,
+                            interactive_,
+                            IsDynamicPartition(install_part.name));
   // Open source fds if we have a delta payload, or for partitions in the
   // partial update.
   bool source_may_exist = manifest_.partial_update() ||
@@ -1407,16 +1407,21 @@
   return true;
 }
 
-bool DeltaPerformer::CheckpointUpdateProgress(bool force) {
+bool DeltaPerformer::ShouldCheckpoint() {
   base::TimeTicks curr_time = base::TimeTicks::Now();
-  if (force || curr_time > update_checkpoint_time_) {
+  if (curr_time > update_checkpoint_time_) {
     update_checkpoint_time_ = curr_time + update_checkpoint_wait_;
-  } else {
+    return true;
+  }
+  return false;
+}
+
+bool DeltaPerformer::CheckpointUpdateProgress(bool force) {
+  if (!force && !ShouldCheckpoint()) {
     return false;
   }
-
   Terminator::set_exit_blocked(true);
-  if (last_updated_buffer_offset_ != buffer_offset_) {
+  if (last_updated_operation_num_ != next_operation_num_) {
     // Resets the progress in case we die in the middle of the state update.
     ResetUpdateProgress(prefs_, true);
     TEST_AND_RETURN_FALSE(prefs_->SetString(
@@ -1426,12 +1431,13 @@
                           signed_hash_calculator_.GetContext()));
     TEST_AND_RETURN_FALSE(
         prefs_->SetInt64(kPrefsUpdateStateNextDataOffset, buffer_offset_));
-    last_updated_buffer_offset_ = buffer_offset_;
+    last_updated_operation_num_ = next_operation_num_;
 
     if (next_operation_num_ < num_total_operations_) {
       size_t partition_index = current_partition_;
-      while (next_operation_num_ >= acc_num_operations_[partition_index])
+      while (next_operation_num_ >= acc_num_operations_[partition_index]) {
         partition_index++;
+      }
       const size_t partition_operation_num =
           next_operation_num_ -
           (partition_index ? acc_num_operations_[partition_index - 1] : 0);
@@ -1524,4 +1530,20 @@
                    part_name) != dynamic_partitions_.end();
 }
 
+std::unique_ptr<PartitionWriter> DeltaPerformer::CreatePartitionWriter(
+    const PartitionUpdate& partition_update,
+    const InstallPlan::Partition& install_part,
+    DynamicPartitionControlInterface* dynamic_control,
+    size_t block_size,
+    bool is_interactive,
+    bool is_dynamic_partition) {
+  return partition_writer::CreatePartitionWriter(
+      partition_update,
+      install_part,
+      dynamic_control,
+      block_size_,
+      interactive_,
+      IsDynamicPartition(install_part.name));
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/delta_performer.h b/payload_consumer/delta_performer.h
index 67ca6ce..ec23997 100644
--- a/payload_consumer/delta_performer.h
+++ b/payload_consumer/delta_performer.h
@@ -202,6 +202,20 @@
       const std::string& update_check_response_hash,
       uint64_t* required_size);
 
+ protected:
+  // Exposed as virtual for testing purposes.
+  virtual std::unique_ptr<PartitionWriter> CreatePartitionWriter(
+      const PartitionUpdate& partition_update,
+      const InstallPlan::Partition& install_part,
+      DynamicPartitionControlInterface* dynamic_control,
+      size_t block_size,
+      bool is_interactive,
+      bool is_dynamic_partition);
+
+  // return true if it has been long enough and a checkpoint should be saved.
+  // Exposed for unittest purposes.
+  virtual bool ShouldCheckpoint();
+
  private:
   friend class DeltaPerformerTest;
   friend class DeltaPerformerIntegrationTest;
@@ -312,6 +326,7 @@
 
   // Check if partition `part_name` is a dynamic partition.
   bool IsDynamicPartition(const std::string& part_name);
+
   // Update Engine preference store.
   PrefsInterface* prefs_;
 
@@ -369,8 +384,8 @@
   // Offset of buffer_ in the binary blobs section of the update.
   uint64_t buffer_offset_{0};
 
-  // Last |buffer_offset_| value updated as part of the progress update.
-  uint64_t last_updated_buffer_offset_{std::numeric_limits<uint64_t>::max()};
+  // Last |next_operation_num_| value updated as part of the progress update.
+  uint64_t last_updated_operation_num_{std::numeric_limits<uint64_t>::max()};
 
   // The block size (parsed from the manifest).
   uint32_t block_size_{0};
diff --git a/payload_consumer/delta_performer_unittest.cc b/payload_consumer/delta_performer_unittest.cc
index ba387a5..840ecf6 100644
--- a/payload_consumer/delta_performer_unittest.cc
+++ b/payload_consumer/delta_performer_unittest.cc
@@ -20,6 +20,8 @@
 #include <inttypes.h>
 #include <time.h>
 
+#include <algorithm>
+#include <map>
 #include <memory>
 #include <string>
 #include <vector>
@@ -31,6 +33,7 @@
 #include <base/strings/string_number_conversions.h>
 #include <base/strings/string_util.h>
 #include <base/strings/stringprintf.h>
+#include <brillo/secure_blob.h>
 #include <gmock/gmock.h>
 #include <google/protobuf/repeated_field.h>
 #include <gtest/gtest.h>
@@ -41,10 +44,12 @@
 #include "update_engine/common/fake_hardware.h"
 #include "update_engine/common/fake_prefs.h"
 #include "update_engine/common/hardware_interface.h"
+#include "update_engine/common/hash_calculator.h"
 #include "update_engine/common/mock_download_action.h"
 #include "update_engine/common/test_utils.h"
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/fake_file_descriptor.h"
+#include "update_engine/payload_consumer/mock_partition_writer.h"
 #include "update_engine/payload_consumer/payload_constants.h"
 #include "update_engine/payload_consumer/payload_metadata.h"
 #include "update_engine/payload_generator/bzip.h"
@@ -60,6 +65,8 @@
 using test_utils::GetBuildArtifactsPath;
 using test_utils::kRandomString;
 using testing::_;
+using testing::Return;
+using ::testing::Sequence;
 
 extern const char* kUnittestPrivateKeyPath;
 extern const char* kUnittestPublicKeyPath;
@@ -274,7 +281,14 @@
                             const string& source_path,
                             bool expect_success) {
     return ApplyPayloadToData(
-        payload_data, source_path, brillo::Blob(), expect_success);
+        &performer_, payload_data, source_path, brillo::Blob(), expect_success);
+  }
+  brillo::Blob ApplyPayloadToData(const brillo::Blob& payload_data,
+                                  const string& source_path,
+                                  const brillo::Blob& target_data,
+                                  bool expect_success) {
+    return ApplyPayloadToData(
+        &performer_, payload_data, source_path, target_data, expect_success);
   }
 
   // Apply the payload provided in |payload_data| reading from the |source_path|
@@ -282,7 +296,8 @@
   // new target file are set to |target_data| before applying the payload.
   // Expect result of performer_.Write() to be |expect_success|.
   // Returns the result of the payload application.
-  brillo::Blob ApplyPayloadToData(const brillo::Blob& payload_data,
+  brillo::Blob ApplyPayloadToData(DeltaPerformer* delta_performer,
+                                  const brillo::Blob& payload_data,
                                   const string& source_path,
                                   const brillo::Blob& target_data,
                                   bool expect_success) {
@@ -302,7 +317,7 @@
         kPartitionNameKernel, install_plan_.source_slot, "/dev/null");
 
     EXPECT_EQ(expect_success,
-              performer_.Write(payload_data.data(), payload_data.size()));
+              delta_performer->Write(payload_data.data(), payload_data.size()));
     EXPECT_EQ(0, performer_.Close());
 
     brillo::Blob partition_data;
@@ -1095,4 +1110,91 @@
   ASSERT_TRUE(DeltaPerformer::CanResumeUpdate(&prefs_, payload_id));
 }
 
+class TestDeltaPerformer : public DeltaPerformer {
+ public:
+  using DeltaPerformer::DeltaPerformer;
+
+  std::unique_ptr<PartitionWriter> CreatePartitionWriter(
+      const PartitionUpdate& partition_update,
+      const InstallPlan::Partition& install_part,
+      DynamicPartitionControlInterface* dynamic_control,
+      size_t block_size,
+      bool is_interactive,
+      bool is_dynamic_partition) {
+    LOG(INFO) << __FUNCTION__ << ": " << install_part.name;
+    auto node = partition_writers_.extract(install_part.name);
+    return std::move(node.mapped());
+  }
+
+  bool ShouldCheckpoint() override { return true; }
+
+  std::map<std::string, std::unique_ptr<MockPartitionWriter>>
+      partition_writers_;
+};
+
+namespace {
+AnnotatedOperation GetSourceCopyOp(uint32_t src_block,
+                                   uint32_t dst_block,
+                                   const void* data,
+                                   size_t length) {
+  AnnotatedOperation aop;
+  *(aop.op.add_src_extents()) = ExtentForRange(0, 1);
+  *(aop.op.add_dst_extents()) = ExtentForRange(0, 1);
+  aop.op.set_type(InstallOperation::SOURCE_COPY);
+  brillo::Blob src_hash;
+  HashCalculator::RawHashOfBytes(data, length, &src_hash);
+  aop.op.set_src_sha256_hash(src_hash.data(), src_hash.size());
+  return aop;
+}
+}  // namespace
+
+TEST_F(DeltaPerformerTest, SetNextOpIndex) {
+  TestDeltaPerformer delta_performer{&prefs_,
+                                     &fake_boot_control_,
+                                     &fake_hardware_,
+                                     &mock_delegate_,
+                                     &install_plan_,
+                                     &payload_,
+                                     false};
+  brillo::Blob expected_data(std::begin(kRandomString),
+                             std::end(kRandomString));
+  expected_data.resize(4096 * 2);  // block size
+  AnnotatedOperation aop;
+
+  ScopedTempFile source("Source-XXXXXX");
+  EXPECT_TRUE(test_utils::WriteFileVector(source.path(), expected_data));
+
+  PartitionConfig old_part(kPartitionNameRoot);
+  old_part.path = source.path();
+  old_part.size = expected_data.size();
+
+  delta_performer.partition_writers_[kPartitionNameRoot] =
+      std::make_unique<MockPartitionWriter>();
+  auto& writer1 = *delta_performer.partition_writers_[kPartitionNameRoot];
+
+  Sequence seq;
+  std::vector<size_t> indices;
+  EXPECT_CALL(writer1, CheckpointUpdateProgress(_))
+      .WillRepeatedly(
+          [&indices](size_t index) mutable { indices.emplace_back(index); });
+  EXPECT_CALL(writer1, Init(_, true, _)).Times(1).WillOnce(Return(true));
+  EXPECT_CALL(writer1, PerformSourceCopyOperation(_, _))
+      .Times(2)
+      .WillRepeatedly(Return(true));
+
+  brillo::Blob payload_data = GeneratePayload(
+      brillo::Blob(),
+      {GetSourceCopyOp(0, 0, expected_data.data(), 4096),
+       GetSourceCopyOp(1, 1, expected_data.data() + 4096, 4096)},
+      false,
+      &old_part);
+
+  ApplyPayloadToData(&delta_performer, payload_data, source.path(), {}, true);
+  ASSERT_TRUE(std::is_sorted(indices.begin(), indices.end()));
+  ASSERT_GT(indices.size(), 0UL);
+
+  // Should be equal to number of operations
+  ASSERT_EQ(indices[indices.size() - 1], 2UL);
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/mock_partition_writer.h b/payload_consumer/mock_partition_writer.h
new file mode 100644
index 0000000..b056010
--- /dev/null
+++ b/payload_consumer/mock_partition_writer.h
@@ -0,0 +1,68 @@
+//
+// Copyright (C) 2020 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#ifndef UPDATE_ENGINE_MOCK_PARTITION_WRITER_H_
+#define UPDATE_ENGINE_MOCK_PARTITION_WRITER_H_
+
+#include <gmock/gmock.h>
+
+#include "common/error_code.h"
+#include "payload_generator/delta_diff_generator.h"
+#include "update_engine/payload_consumer/partition_writer.h"
+
+namespace chromeos_update_engine {
+class MockPartitionWriter : public PartitionWriter {
+ public:
+  MockPartitionWriter() : PartitionWriter({}, {}, nullptr, kBlockSize, false) {}
+  virtual ~MockPartitionWriter() = default;
+
+  // Perform necessary initialization work before InstallOperation can be
+  // applied to this partition
+  MOCK_METHOD(bool, Init, (const InstallPlan*, bool, size_t), (override));
+
+  // |CheckpointUpdateProgress| will be called after SetNextOpIndex(), but it's
+  // optional. DeltaPerformer may or may not call this everytime an operation is
+  // applied.
+  MOCK_METHOD(void, CheckpointUpdateProgress, (size_t), (override));
+
+  // These perform a specific type of operation and return true on success.
+  // |error| will be set if source hash mismatch, otherwise |error| might not be
+  // set even if it fails.
+  MOCK_METHOD(bool,
+              PerformReplaceOperation,
+              (const InstallOperation&, const void*, size_t),
+              (override));
+  MOCK_METHOD(bool,
+              PerformZeroOrDiscardOperation,
+              (const InstallOperation&),
+              (override));
+  MOCK_METHOD(bool,
+              PerformSourceCopyOperation,
+              (const InstallOperation&, ErrorCode*),
+              (override));
+  MOCK_METHOD(bool,
+              PerformSourceBsdiffOperation,
+              (const InstallOperation&, ErrorCode*, const void*, size_t),
+              (override));
+  MOCK_METHOD(bool,
+              PerformPuffDiffOperation,
+              (const InstallOperation&, ErrorCode*, const void*, size_t),
+              (override));
+};
+
+}  // namespace chromeos_update_engine
+
+#endif
diff --git a/payload_consumer/partition_writer.h b/payload_consumer/partition_writer.h
index 4b420d2..82e557a 100644
--- a/payload_consumer/partition_writer.h
+++ b/payload_consumer/partition_writer.h
@@ -29,6 +29,7 @@
 #include "update_engine/payload_consumer/file_descriptor.h"
 #include "update_engine/payload_consumer/install_plan.h"
 #include "update_engine/update_metadata.pb.h"
+
 namespace chromeos_update_engine {
 class PartitionWriter {
  public: