Refactor for ISnapshotWriter removal.

Bug: 280529365
Test: update_engine_unittests
      apply ota
Change-Id: I03e899ba31077f828268be3775ebceba3850c528
diff --git a/Android.bp b/Android.bp
index 144dd30..a33aa8f 100644
--- a/Android.bp
+++ b/Android.bp
@@ -1247,6 +1247,7 @@
     host_supported: true,
     recovery_available: true,
     ramdisk_available: true,
+    vendor_ramdisk_available: true,
 
     target: {
         darwin: {
diff --git a/aosp/dynamic_partition_control_android.cc b/aosp/dynamic_partition_control_android.cc
index 19ea6ec..a309b6e 100644
--- a/aosp/dynamic_partition_control_android.cc
+++ b/aosp/dynamic_partition_control_android.cc
@@ -1424,11 +1424,11 @@
   return metadata_device_ != nullptr;
 }
 
-std::unique_ptr<android::snapshot::ISnapshotWriter>
+std::unique_ptr<android::snapshot::ICowWriter>
 DynamicPartitionControlAndroid::OpenCowWriter(
     const std::string& partition_name,
     const std::optional<std::string>& source_path,
-    bool) {
+    std::optional<uint64_t> label) {
   auto suffix = SlotSuffixForSlotNumber(target_slot_);
 
   auto super_device = GetSuperDevice();
@@ -1443,29 +1443,26 @@
       .timeout_ms = kMapSnapshotTimeout};
   // TODO(zhangkelvin) Open an APPEND mode CowWriter once there's an API to do
   // it.
-  return snapshot_->OpenSnapshotWriter(params, std::move(source_path));
+  return snapshot_->OpenSnapshotWriter(params, label);
 }  // namespace chromeos_update_engine
 
 std::unique_ptr<FileDescriptor> DynamicPartitionControlAndroid::OpenCowFd(
     const std::string& unsuffixed_partition_name,
     const std::optional<std::string>& source_path,
     bool is_append) {
-  auto cow_writer =
-      OpenCowWriter(unsuffixed_partition_name, source_path, is_append);
+  auto cow_writer = OpenCowWriter(
+      unsuffixed_partition_name, source_path, {kEndOfInstallLabel});
   if (cow_writer == nullptr) {
+    LOG(ERROR) << "OpenCowWriter failed";
     return nullptr;
   }
-  if (!cow_writer->InitializeAppend(kEndOfInstallLabel)) {
-    LOG(ERROR) << "Failed to InitializeAppend(" << kEndOfInstallLabel << ")";
+  auto fd = cow_writer->OpenFileDescriptor(source_path);
+  if (fd == nullptr) {
+    LOG(ERROR) << "ICowReader::OpenFileDescriptor failed";
     return nullptr;
   }
-  auto reader = cow_writer->OpenReader();
-  if (reader == nullptr) {
-    LOG(ERROR) << "ICowWriter::OpenReader() failed.";
-    return nullptr;
-  }
-  return std::make_unique<CowWriterFileDescriptor>(std::move(cow_writer),
-                                                   std::move(reader));
+  return std::make_unique<CowWriterFileDescriptor>(
+      std::move(cow_writer), std::move(fd), source_path);
 }
 
 std::optional<base::FilePath> DynamicPartitionControlAndroid::GetSuperDevice() {
diff --git a/aosp/dynamic_partition_control_android.h b/aosp/dynamic_partition_control_android.h
index 9851a99..830901b 100644
--- a/aosp/dynamic_partition_control_android.h
+++ b/aosp/dynamic_partition_control_android.h
@@ -17,16 +17,16 @@
 #ifndef UPDATE_ENGINE_AOSP_DYNAMIC_PARTITION_CONTROL_ANDROID_H_
 #define UPDATE_ENGINE_AOSP_DYNAMIC_PARTITION_CONTROL_ANDROID_H_
 
+#include <array>
 #include <memory>
 #include <set>
 #include <string>
 #include <string_view>
-#include <array>
+#include <vector>
 
 #include <base/files/file_util.h>
 #include <libsnapshot/auto_device.h>
 #include <libsnapshot/snapshot.h>
-#include <libsnapshot/snapshot_writer.h>
 
 #include "update_engine/common/dynamic_partition_control_interface.h"
 
@@ -102,14 +102,14 @@
 
   // Partition name is expected to be unsuffixed. e.g. system, vendor
   // Return an interface to write to a snapshoted partition.
-  std::unique_ptr<android::snapshot::ISnapshotWriter> OpenCowWriter(
+  std::unique_ptr<android::snapshot::ICowWriter> OpenCowWriter(
       const std::string& unsuffixed_partition_name,
       const std::optional<std::string>& source_path,
-      bool is_append) override;
+      std::optional<uint64_t> label) override;
   std::unique_ptr<FileDescriptor> OpenCowFd(
       const std::string& unsuffixed_partition_name,
       const std::optional<std::string>&,
-      bool is_append = false) override;
+      bool is_append) override;
 
   bool MapAllPartitions() override;
   bool UnmapAllPartitions() override;
diff --git a/aosp/mock_dynamic_partition_control_android.h b/aosp/mock_dynamic_partition_control_android.h
index e940e50..cc6ebf3 100644
--- a/aosp/mock_dynamic_partition_control_android.h
+++ b/aosp/mock_dynamic_partition_control_android.h
@@ -17,13 +17,13 @@
 #include <stdint.h>
 
 #include <memory>
+#include <optional>
 #include <set>
 #include <string>
 
 #include <gmock/gmock.h>
 
 #include <libsnapshot/cow_writer.h>
-#include <libsnapshot/snapshot_writer.h>
 
 #include "payload_consumer/file_descriptor.h"
 #include "update_engine/aosp/dynamic_partition_control_android.h"
@@ -96,11 +96,11 @@
               PrepareDynamicPartitionsForUpdate,
               (uint32_t, uint32_t, const DeltaArchiveManifest&, bool),
               (override));
-  MOCK_METHOD(std::unique_ptr<android::snapshot::ISnapshotWriter>,
+  MOCK_METHOD(std::unique_ptr<android::snapshot::ICowWriter>,
               OpenCowWriter,
               (const std::string& unsuffixed_partition_name,
                const std::optional<std::string>& source_path,
-               bool is_append),
+               std::optional<uint64_t> label),
               (override));
   MOCK_METHOD(std::unique_ptr<FileDescriptor>,
               OpenCowFd,
diff --git a/common/dynamic_partition_control_interface.h b/common/dynamic_partition_control_interface.h
index 3f735bb..490892e 100644
--- a/common/dynamic_partition_control_interface.h
+++ b/common/dynamic_partition_control_interface.h
@@ -32,7 +32,7 @@
 
 // Forware declare for libsnapshot/snapshot_writer.h
 namespace android::snapshot {
-class ISnapshotWriter;
+class ICowWriter;
 }
 
 namespace chromeos_update_engine {
@@ -166,10 +166,10 @@
   // If `is_append` is false, then existing COW data will be overwritten.
   // Otherwise the cow writer will be opened on APPEND mode, existing COW data
   // is preserved.
-  virtual std::unique_ptr<android::snapshot::ISnapshotWriter> OpenCowWriter(
+  virtual std::unique_ptr<android::snapshot::ICowWriter> OpenCowWriter(
       const std::string& unsuffixed_partition_name,
       const std::optional<std::string>&,
-      bool is_append = false) = 0;
+      std::optional<uint64_t> label) = 0;
   // Open a general purpose FD capable to reading and writing to COW. Note that
   // writes must be block aligned.
   virtual std::unique_ptr<FileDescriptor> OpenCowFd(
diff --git a/common/dynamic_partition_control_stub.cc b/common/dynamic_partition_control_stub.cc
index 03f7361..fd9a3b4 100644
--- a/common/dynamic_partition_control_stub.cc
+++ b/common/dynamic_partition_control_stub.cc
@@ -100,11 +100,11 @@
   return true;
 }
 
-std::unique_ptr<android::snapshot::ISnapshotWriter>
+std::unique_ptr<android::snapshot::ICowWriter>
 DynamicPartitionControlStub::OpenCowWriter(
     const std::string& /*unsuffixed_partition_name*/,
     const std::optional<std::string>& /*source_path*/,
-    bool /*is_append*/) {
+    std::optional<uint64_t>) {
   return nullptr;
 }
 
diff --git a/common/dynamic_partition_control_stub.h b/common/dynamic_partition_control_stub.h
index 4236051..1db6a78 100644
--- a/common/dynamic_partition_control_stub.h
+++ b/common/dynamic_partition_control_stub.h
@@ -27,7 +27,8 @@
 
 namespace chromeos_update_engine {
 
-class DynamicPartitionControlStub final : public DynamicPartitionControlInterface {
+class DynamicPartitionControlStub final
+    : public DynamicPartitionControlInterface {
  public:
   FeatureFlag GetDynamicPartitionsFeatureFlag() override;
   FeatureFlag GetVirtualAbFeatureFlag() override;
@@ -62,10 +63,10 @@
       uint32_t target_slot,
       const std::vector<std::string>& partitions) override;
 
-  std::unique_ptr<android::snapshot::ISnapshotWriter> OpenCowWriter(
+  std::unique_ptr<android::snapshot::ICowWriter> OpenCowWriter(
       const std::string& unsuffixed_partition_name,
       const std::optional<std::string>&,
-      bool is_append) override;
+      std::optional<uint64_t> label) override;
 
   std::unique_ptr<FileDescriptor> OpenCowFd(
       const std::string& unsuffixed_partition_name,
diff --git a/common/mock_dynamic_partition_control.h b/common/mock_dynamic_partition_control.h
index 8cff28f..79909b4 100644
--- a/common/mock_dynamic_partition_control.h
+++ b/common/mock_dynamic_partition_control.h
@@ -22,6 +22,7 @@
 #include <vector>
 
 #include <gmock/gmock.h>
+#include <libsnapshot/cow_writer.h>
 
 #include "update_engine/common/dynamic_partition_control_interface.h"
 #include "update_engine/payload_consumer/file_descriptor.h"
@@ -58,9 +59,11 @@
               (const std::string&, const InstallOperation&, InstallOperation*),
               (override));
 
-  MOCK_METHOD(std::unique_ptr<android::snapshot::ISnapshotWriter>,
+  MOCK_METHOD(std::unique_ptr<android::snapshot::ICowWriter>,
               OpenCowWriter,
-              (const std::string&, const std::optional<std::string>&, bool),
+              (const std::string&,
+               const std::optional<std::string>&,
+               std::optional<uint64_t> label),
               (override));
 
   MOCK_METHOD(
diff --git a/payload_consumer/cow_writer_file_descriptor.cc b/payload_consumer/cow_writer_file_descriptor.cc
index 2dca1b3..a7d43a7 100644
--- a/payload_consumer/cow_writer_file_descriptor.cc
+++ b/payload_consumer/cow_writer_file_descriptor.cc
@@ -17,6 +17,7 @@
 #include "update_engine/payload_consumer/cow_writer_file_descriptor.h"
 
 #include <memory>
+#include <string>
 #include <utility>
 
 #include <base/logging.h>
@@ -25,18 +26,14 @@
 #include "update_engine/payload_consumer/file_descriptor.h"
 
 namespace chromeos_update_engine {
-CowWriterFileDescriptor::CowWriterFileDescriptor(
-    std::unique_ptr<android::snapshot::ISnapshotWriter> cow_writer)
-    : cow_writer_(std::move(cow_writer)),
-      cow_reader_(cow_writer_->OpenReader()) {
-  CHECK_NE(cow_writer_, nullptr);
-  CHECK_NE(cow_reader_, nullptr);
-}
 
 CowWriterFileDescriptor::CowWriterFileDescriptor(
-    std::unique_ptr<android::snapshot::ISnapshotWriter> cow_writer,
-    std::unique_ptr<FileDescriptor> cow_reader)
-    : cow_writer_(std::move(cow_writer)), cow_reader_(std::move(cow_reader)) {
+    std::unique_ptr<android::snapshot::ICowWriter> cow_writer,
+    std::unique_ptr<FileDescriptor> cow_reader,
+    const std::optional<std::string>& source_device)
+    : cow_writer_(std::move(cow_writer)),
+      cow_reader_(std::move(cow_reader)),
+      source_device_(source_device) {
   CHECK_NE(cow_writer_, nullptr);
   CHECK_NE(cow_reader_, nullptr);
 }
@@ -68,9 +65,10 @@
       LOG(ERROR) << "Failed to Finalize() cow writer";
       return -1;
     }
-    cow_reader_ = cow_writer_->OpenReader();
+    cow_reader_ = cow_writer_->OpenFileDescriptor(source_device_);
     if (cow_reader_ == nullptr) {
-      LOG(ERROR) << "Failed to re-open cow reader after writing to COW";
+      LOG(ERROR)
+          << "Failed to re-open cow file descriptor after writing to COW";
       return -1;
     }
     const auto pos = cow_reader_->Seek(offset, SEEK_SET);
diff --git a/payload_consumer/cow_writer_file_descriptor.h b/payload_consumer/cow_writer_file_descriptor.h
index fbf0a0d..24fef83 100644
--- a/payload_consumer/cow_writer_file_descriptor.h
+++ b/payload_consumer/cow_writer_file_descriptor.h
@@ -16,8 +16,9 @@
 
 #include <cstdint>
 #include <memory>
+#include <string>
 
-#include <libsnapshot/snapshot_writer.h>
+#include <libsnapshot/cow_writer.h>
 
 #include "update_engine/payload_consumer/file_descriptor.h"
 
@@ -28,13 +29,12 @@
 // FEC. Writes must be block aligned(4096) or write will fail.
 class CowWriterFileDescriptor final : public FileDescriptor {
  public:
-  explicit CowWriterFileDescriptor(
-      std::unique_ptr<android::snapshot::ISnapshotWriter> cow_writer);
-
-  // |cow_reader| should be obtained by calling |cow_writer->OpenReader()|
+  // |cow_reader| should be obtained by calling
+  // |cow_writer->OpenReader()->OpenFileDescriptor()|
   CowWriterFileDescriptor(
-      std::unique_ptr<android::snapshot::ISnapshotWriter> cow_writer,
-      std::unique_ptr<FileDescriptor> cow_reader);
+      std::unique_ptr<android::snapshot::ICowWriter> cow_writer,
+      std::unique_ptr<FileDescriptor> cow_reader,
+      const std::optional<std::string>& source_device);
   ~CowWriterFileDescriptor();
 
   bool Open(const char* path, int flags, mode_t mode) override;
@@ -64,8 +64,9 @@
   bool IsOpen() override;
 
  private:
-  std::unique_ptr<android::snapshot::ISnapshotWriter> cow_writer_;
+  std::unique_ptr<android::snapshot::ICowWriter> cow_writer_;
   FileDescriptorPtr cow_reader_;
+  std::optional<std::string> source_device_;
   bool dirty_ = false;
 };
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/cow_writer_file_descriptor_unittest.cc b/payload_consumer/cow_writer_file_descriptor_unittest.cc
index c596e3b..a39f13f 100644
--- a/payload_consumer/cow_writer_file_descriptor_unittest.cc
+++ b/payload_consumer/cow_writer_file_descriptor_unittest.cc
@@ -18,13 +18,14 @@
 
 #include <cstring>
 #include <memory>
+#include <string>
 #include <utility>
 #include <vector>
 
 #include <android-base/unique_fd.h>
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include <libsnapshot/snapshot_writer.h>
+#include <libsnapshot/cow_writer.h>
 
 #include "update_engine/common/utils.h"
 
@@ -33,9 +34,8 @@
 constexpr size_t PARTITION_SIZE = BLOCK_SIZE * 10;
 
 using android::base::unique_fd;
-using android::snapshot::CompressedSnapshotWriter;
 using android::snapshot::CowOptions;
-using android::snapshot::ISnapshotWriter;
+using android::snapshot::ICowWriter;
 
 class CowWriterFileDescriptorUnittest : public ::testing::Test {
  public:
@@ -48,18 +48,22 @@
         << strerror(errno);
   }
 
-  std::unique_ptr<CompressedSnapshotWriter> GetCowWriter() {
+  std::unique_ptr<ICowWriter> GetCowWriter() {
+    static constexpr uint32_t kTestCowVersion = 2;
     const CowOptions options{.block_size = BLOCK_SIZE, .compression = "gz"};
-    auto snapshot_writer = std::make_unique<CompressedSnapshotWriter>(options);
     int fd = open(cow_device_file_.path().c_str(), O_RDWR);
     EXPECT_NE(fd, -1);
-    EXPECT_TRUE(snapshot_writer->SetCowDevice(unique_fd{fd}));
-    snapshot_writer->SetSourceDevice(cow_source_file_.path());
-    return snapshot_writer;
+    return android::snapshot::CreateCowWriter(
+        kTestCowVersion, options, unique_fd{fd});
   }
-  CowWriterFileDescriptor GetCowFd() {
+  std::unique_ptr<CowWriterFileDescriptor> GetCowFd() {
     auto cow_writer = GetCowWriter();
-    return CowWriterFileDescriptor{std::move(cow_writer)};
+    EXPECT_NE(cow_writer, nullptr);
+    auto fd = cow_writer->OpenFileDescriptor({cow_source_file_.path()});
+    EXPECT_NE(fd, nullptr);
+    auto source_path = std::optional<std::string>{cow_source_file_.path()};
+    return std::make_unique<CowWriterFileDescriptor>(
+        std::move(cow_writer), std::move(fd), source_path);
   }
 
   ScopedTempFile cow_source_file_{"cow_source.XXXXXX", true};
@@ -76,7 +80,6 @@
   std::fill(verity_data.begin(), verity_data.end(), 0xAA);
 
   auto cow_writer = GetCowWriter();
-  cow_writer->Initialize();
 
   // Simulate Writing InstallOp data
   ASSERT_TRUE(cow_writer->AddRawBlocks(0, buffer.data(), buffer.size()));
@@ -88,11 +91,7 @@
       cow_writer->AddRawBlocks(4, verity_data.data(), verity_data.size()));
   ASSERT_TRUE(cow_writer->Finalize());
 
-  cow_writer = GetCowWriter();
-  ASSERT_NE(nullptr, cow_writer);
-  ASSERT_TRUE(cow_writer->InitializeAppend(23));
-  auto cow_fd =
-      std::make_unique<CowWriterFileDescriptor>(std::move(cow_writer));
+  auto cow_fd = GetCowFd();
 
   ASSERT_EQ((ssize_t)BLOCK_SIZE * 4, cow_fd->Seek(BLOCK_SIZE * 4, SEEK_SET));
   std::vector<unsigned char> read_back(4096);
@@ -103,16 +102,13 @@
   // Since we didn't write anything to this instance of cow_fd, destructor
   // should not call Finalize(). As finalize will drop ops after resume label,
   // causing subsequent reads to fail.
-  cow_writer = GetCowWriter();
-  ASSERT_NE(nullptr, cow_writer);
-  ASSERT_TRUE(cow_writer->InitializeAppend(23));
-  cow_fd = std::make_unique<CowWriterFileDescriptor>(std::move(cow_writer));
+  cow_fd = GetCowFd();
 
   ASSERT_EQ((ssize_t)BLOCK_SIZE * 4, cow_fd->Seek(BLOCK_SIZE * 4, SEEK_SET));
   ASSERT_EQ((ssize_t)read_back.size(),
             cow_fd->Read(read_back.data(), read_back.size()));
   ASSERT_EQ(verity_data, read_back)
-      << "Could not read verity data afeter InitializeAppend() => Read() => "
+      << "Could not read verity data after InitializeAppend() => Read() => "
          "InitializeAppend() sequence. If no writes happened while CowWriterFd "
          "is open, Finalize() should not be called.";
 }
diff --git a/payload_consumer/file_descriptor.h b/payload_consumer/file_descriptor.h
index f672871..3f0361a 100644
--- a/payload_consumer/file_descriptor.h
+++ b/payload_consumer/file_descriptor.h
@@ -21,7 +21,7 @@
 #include <sys/types.h>
 #include <memory>
 
-#include <base/macros.h>
+#include <android-base/macros.h>
 
 // Abstraction for managing opening, reading, writing and closing of file
 // descriptors. This includes an abstract class and one standard implementation
diff --git a/payload_consumer/filesystem_verifier_action_unittest.cc b/payload_consumer/filesystem_verifier_action_unittest.cc
index 3c3997d..565976a 100644
--- a/payload_consumer/filesystem_verifier_action_unittest.cc
+++ b/payload_consumer/filesystem_verifier_action_unittest.cc
@@ -29,7 +29,7 @@
 #include <brillo/secure_blob.h>
 #include <fec/ecc.h>
 #include <gtest/gtest.h>
-#include <libsnapshot/snapshot_writer.h>
+#include <libsnapshot/cow_writer.h>
 #include <sys/stat.h>
 
 #include "gmock/gmock-spec-builders.h"
diff --git a/payload_consumer/snapshot_extent_writer_unittest.cc b/payload_consumer/snapshot_extent_writer_unittest.cc
index 01c3d2e..057dbda 100644
--- a/payload_consumer/snapshot_extent_writer_unittest.cc
+++ b/payload_consumer/snapshot_extent_writer_unittest.cc
@@ -17,7 +17,9 @@
 #include <array>
 #include <cstring>
 #include <map>
+#include <memory>
 #include <numeric>
+#include <string>
 #include <vector>
 
 #include <gtest/gtest.h>
@@ -89,6 +91,14 @@
   uint32_t GetBlockSize() const override { return 4096; }
   std::optional<uint32_t> GetMaxBlocks() const override { return {}; }
 
+  std::unique_ptr<android::snapshot::ICowReader> OpenReader() override {
+    return nullptr;
+  }
+  std::unique_ptr<FileDescriptor> OpenFileDescriptor(
+      const std::optional<std::string>&) override {
+    return nullptr;
+  }
+
   // Return number of bytes the cow image occupies on disk.
   uint64_t GetCowSize() override {
     return std::accumulate(
diff --git a/payload_consumer/vabc_partition_writer.cc b/payload_consumer/vabc_partition_writer.cc
index b00ff70..bcaccce 100644
--- a/payload_consumer/vabc_partition_writer.cc
+++ b/payload_consumer/vabc_partition_writer.cc
@@ -157,22 +157,25 @@
     // TODO(zhangkelvin) Make |source_path| a std::optional<std::string>
     source_path = install_part_.source_path;
   }
-  cow_writer_ = dynamic_control_->OpenCowWriter(
-      install_part_.name, source_path, install_plan->is_resume);
-  TEST_AND_RETURN_FALSE(cow_writer_ != nullptr);
 
   // ===== Resume case handling code goes here ====
   // It is possible that the SOURCE_COPY are already written but
   // |next_op_index_| is still 0. In this case we discard previously written
   // SOURCE_COPY, and start over.
+  std::optional<uint64_t> label;
   if (install_plan->is_resume && next_op_index > 0) {
     LOG(INFO) << "Resuming update on partition `"
               << partition_update_.partition_name() << "` op index "
               << next_op_index;
-    TEST_AND_RETURN_FALSE(cow_writer_->InitializeAppend(next_op_index));
+    label = {next_op_index};
+  }
+
+  cow_writer_ =
+      dynamic_control_->OpenCowWriter(install_part_.name, source_path, label);
+  TEST_AND_RETURN_FALSE(cow_writer_ != nullptr);
+
+  if (label) {
     return true;
-  } else {
-    TEST_AND_RETURN_FALSE(cow_writer_->Initialize());
   }
 
   // ==============================================
@@ -383,7 +386,10 @@
   TEST_AND_RETURN_FALSE(cow_writer_ != nullptr);
   TEST_AND_RETURN_FALSE(cow_writer_->AddLabel(kEndOfInstallLabel));
   TEST_AND_RETURN_FALSE(cow_writer_->Finalize());
-  TEST_AND_RETURN_FALSE(cow_writer_->VerifyMergeOps());
+
+  auto cow_reader = cow_writer_->OpenReader();
+  TEST_AND_RETURN_FALSE(cow_reader);
+  TEST_AND_RETURN_FALSE(cow_reader->VerifyMergeOps());
   return true;
 }
 
diff --git a/payload_consumer/vabc_partition_writer.h b/payload_consumer/vabc_partition_writer.h
index 889f376..977fbe5 100644
--- a/payload_consumer/vabc_partition_writer.h
+++ b/payload_consumer/vabc_partition_writer.h
@@ -22,7 +22,7 @@
 #include <string>
 #include <vector>
 
-#include <libsnapshot/snapshot_writer.h>
+#include <libsnapshot/cow_writer.h>
 
 #include "update_engine/payload_consumer/extent_map.h"
 #include "update_engine/payload_consumer/install_operation_executor.h"
@@ -73,7 +73,7 @@
   [[nodiscard]] bool DoesDeviceSupportsXor();
   bool IsXorEnabled() const noexcept { return xor_map_.size() > 0; }
   [[nodiscard]] bool WriteAllCopyOps();
-  std::unique_ptr<android::snapshot::ISnapshotWriter> cow_writer_;
+  std::unique_ptr<android::snapshot::ICowWriter> cow_writer_;
 
   [[nodiscard]] std::unique_ptr<ExtentWriter> CreateBaseExtentWriter();
 
diff --git a/payload_consumer/vabc_partition_writer_unittest.cc b/payload_consumer/vabc_partition_writer_unittest.cc
index 93c692b..33a70f3 100644
--- a/payload_consumer/vabc_partition_writer_unittest.cc
+++ b/payload_consumer/vabc_partition_writer_unittest.cc
@@ -22,7 +22,7 @@
 #include <bsdiff/bsdiff.h>
 #include <gtest/gtest.h>
 #include <libsnapshot/cow_writer.h>
-#include <libsnapshot/mock_snapshot_writer.h>
+#include <libsnapshot/mock_cow_writer.h>
 
 #include "update_engine/common/error_code.h"
 #include "update_engine/common/hash_calculator.h"
@@ -75,7 +75,7 @@
 
   android::snapshot::CowOptions options_ = {
       .block_size = static_cast<uint32_t>(kBlockSize)};
-  android::snapshot::MockSnapshotWriter cow_writer_;
+  android::snapshot::MockCowWriter cow_writer_;
   MockDynamicPartitionControl dynamic_control_;
   PartitionUpdate partition_update_;
   InstallPlan install_plan_;
@@ -93,14 +93,12 @@
   AddMergeOp(&partition_update_, {42, 5}, {40, 5}, CowMergeOperation::COW_XOR);
   VABCPartitionWriter writer_{
       partition_update_, install_part_, &dynamic_control_, kBlockSize};
-  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, false))
+  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, _))
       .WillOnce(Invoke([](const std::string&,
                           const std::optional<std::string>&,
-                          bool) {
-        auto cow_writer =
-            std::make_unique<android::snapshot::MockSnapshotWriter>();
+                          std::optional<uint64_t>) {
+        auto cow_writer = std::make_unique<android::snapshot::MockCowWriter>();
         auto expected_merge_sequence = {10, 14, 13, 20, 25, 40, 41, 42, 43, 44};
-        EXPECT_CALL(*cow_writer, Initialize()).WillOnce(Return(true));
         EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
             .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
             .WillOnce(Return(true));
@@ -118,13 +116,14 @@
       ->set_src_offset(1);
   VABCPartitionWriter writer_{
       partition_update_, install_part_, &dynamic_control_, kBlockSize};
-  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, false))
-      .WillOnce(Invoke(
-          [](const std::string&, const std::optional<std::string>&, bool) {
+  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, _))
+      .WillOnce(
+          Invoke([](const std::string&,
+                    const std::optional<std::string>&,
+                    std::optional<uint64_t>) {
             auto cow_writer =
-                std::make_unique<android::snapshot::MockSnapshotWriter>();
+                std::make_unique<android::snapshot::MockCowWriter>();
             auto expected_merge_sequence = {19, 20, 21};
-            EXPECT_CALL(*cow_writer, Initialize()).WillOnce(Return(true));
             EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
                 .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
                 .WillOnce(Return(true));
@@ -170,16 +169,13 @@
   AddMergeOp(&partition_update_, {20, 1}, {25, 1}, CowMergeOperation::COW_COPY);
   VABCPartitionWriter writer_{
       partition_update_, install_part_, &dynamic_control_, kBlockSize};
-  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, false))
+  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, _))
       .WillOnce(Invoke([xor_enabled](const std::string&,
                                      const std::optional<std::string>&,
-                                     bool) {
-        auto cow_writer =
-            std::make_unique<android::snapshot::MockSnapshotWriter>();
+                                     std::optional<uint64_t>) {
+        auto cow_writer = std::make_unique<android::snapshot::MockCowWriter>();
         ON_CALL(*cow_writer, AddCopy(_, _, _)).WillByDefault(Return(true));
         ON_CALL(*cow_writer, AddLabel(_)).WillByDefault(Return(true));
-        ON_CALL(*cow_writer, Initialize()).WillByDefault(Return(true));
-        EXPECT_CALL(*cow_writer, Initialize());
         if (xor_enabled) {
           EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
               .WillOnce(Return(true));
@@ -246,18 +242,16 @@
   data_hash->assign(reinterpret_cast<const char*>(expected_hash.data()),
                     expected_hash.size());
 
-  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, false))
+  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, _))
       .WillOnce(Invoke([](const std::string&,
                           const std::optional<std::string>&,
-                          bool) {
-        auto cow_writer =
-            std::make_unique<android::snapshot::MockSnapshotWriter>();
+                          std::optional<uint64_t>) {
+        auto cow_writer = std::make_unique<android::snapshot::MockCowWriter>();
         ON_CALL(*cow_writer, AddLabel(_)).WillByDefault(Return(true));
         auto expected_merge_sequence = {10, 11, 13, 14};
         auto expected_merge_sequence_rev = {11, 10, 14, 13};
         const bool is_ascending = android::base::GetBoolProperty(
             "ro.virtual_ab.userspace.snapshots.enabled", false);
-        ON_CALL(*cow_writer, Initialize()).WillByDefault(Return(true));
         if (!is_ascending) {
           EXPECT_CALL(*cow_writer, AddSequenceData(_, _))
               .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence_rev)))
@@ -267,7 +261,6 @@
               .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
               .WillOnce(Return(true));
         }
-        EXPECT_CALL(*cow_writer, Initialize()).Times(1);
         EXPECT_CALL(*cow_writer, AddCopy(_, _, _)).Times(0);
         EXPECT_CALL(*cow_writer, AddRawBlocks(_, _, _)).WillOnce(Return(true));
         EXPECT_CALL(*cow_writer, AddXorBlocks(10, _, kBlockSize * 2, 5, 0))
diff --git a/payload_consumer/xor_extent_writer_unittest.cc b/payload_consumer/xor_extent_writer_unittest.cc
index e6cd89a..45796a6 100644
--- a/payload_consumer/xor_extent_writer_unittest.cc
+++ b/payload_consumer/xor_extent_writer_unittest.cc
@@ -20,7 +20,7 @@
 
 #include <android-base/file.h>
 #include <gtest/gtest.h>
-#include <libsnapshot/mock_snapshot_writer.h>
+#include <libsnapshot/mock_cow_writer.h>
 
 #include "common/utils.h"
 #include "update_engine/payload_consumer/extent_map.h"
@@ -60,7 +60,7 @@
   InstallOperation op_;
   FileDescriptorPtr source_fd_ = std::make_shared<EintrSafeFileDescriptor>();
   ExtentMap<const CowMergeOperation*> xor_map_;
-  android::snapshot::MockSnapshotWriter cow_writer_;
+  android::snapshot::MockCowWriter cow_writer_;
   TemporaryFile source_part_;
   TemporaryFile target_part_;
 };