15 files changed, 1382 insertions, 460 deletions

diff --git a/tests/Android.mk b/tests/Android.mk
index fe4de31a2..f19ca36c0 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -18,7 +18,7 @@ LOCAL_PATH := $(call my-dir)
 
 # Unit tests
 include $(CLEAR_VARS)
-LOCAL_CFLAGS := -Werror
+LOCAL_CFLAGS := -Wall -Werror
 LOCAL_MODULE := recovery_unit_test
 LOCAL_COMPATIBILITY_SUITE := device-tests
 LOCAL_STATIC_LIBRARIES := \
@@ -30,7 +30,8 @@ LOCAL_STATIC_LIBRARIES := \
     libutils \
     libz \
     libselinux \
-    libbase
+    libbase \
+    libBionicGtestMain
 
 LOCAL_SRC_FILES := \
     unit/asn1_decoder_test.cpp \
@@ -47,11 +48,12 @@ include $(BUILD_NATIVE_TEST)
 
 # Manual tests
 include $(CLEAR_VARS)
-LOCAL_CFLAGS := -Werror
+LOCAL_CFLAGS := -Wall -Werror
 LOCAL_MODULE := recovery_manual_test
 LOCAL_STATIC_LIBRARIES := \
     libminui \
-    libbase
+    libbase \
+    libBionicGtestMain
 
 LOCAL_SRC_FILES := manual/recovery_test.cpp
 LOCAL_SHARED_LIBRARIES := \
@@ -82,6 +84,7 @@ include $(BUILD_NATIVE_TEST)
 # Component tests
 include $(CLEAR_VARS)
 LOCAL_CFLAGS := \
+    -Wall \
     -Werror \
     -D_FILE_OFFSET_BITS=64
 
@@ -143,9 +146,9 @@ LOCAL_STATIC_LIBRARIES := \
     libdivsufsort \
     libdivsufsort64 \
     libfs_mgr \
-    liblog \
     libvintf_recovery \
     libvintf \
+    libhidl-gen-utils \
     libtinyxml2 \
     libselinux \
     libext4_utils \
@@ -154,6 +157,7 @@ LOCAL_STATIC_LIBRARIES := \
     libcrypto \
     libbz \
     libziparchive \
+    liblog \
     libutils \
     libz \
     libbase \
@@ -163,6 +167,7 @@ LOCAL_STATIC_LIBRARIES := \
     libsquashfs_utils \
     libcutils \
     libbrotli \
+    libBionicGtestMain \
     $(tune2fs_static_libraries)
 
 testdata_files := $(call find-subdir-files, testdata/*)
@@ -192,7 +197,7 @@ include $(BUILD_NATIVE_TEST)
 
 # Host tests
 include $(CLEAR_VARS)
-LOCAL_CFLAGS := -Werror
+LOCAL_CFLAGS := -Wall -Werror
 LOCAL_MODULE := recovery_host_test
 LOCAL_MODULE_HOST_OS := linux
 LOCAL_C_INCLUDES := bootable/recovery
@@ -201,16 +206,19 @@ LOCAL_SRC_FILES := \
 LOCAL_STATIC_LIBRARIES := \
     libimgdiff \
     libimgpatch \
+    libotautil \
     libbsdiff \
     libbspatch \
     libziparchive \
     libutils \
     libbase \
     libcrypto \
+    libbrotli \
     libbz \
     libdivsufsort64 \
     libdivsufsort \
-    libz
+    libz \
+    libBionicGtestMain
 LOCAL_SHARED_LIBRARIES := \
     liblog
 include $(BUILD_HOST_NATIVE_TEST)
diff --git a/tests/common/test_constants.h b/tests/common/test_constants.h
index f6b6922a4..514818e0a 100644
--- a/tests/common/test_constants.h
+++ b/tests/common/test_constants.h
@@ -19,6 +19,8 @@
 
 #include <stdlib.h>
 
+#include <string>
+
 // Zip entries in ziptest_valid.zip.
 static const std::string kATxtContents("abcdefghabcdefgh\n");
 static const std::string kBTxtContents("abcdefgh\n");
@@ -30,10 +32,14 @@ static const std::string kATxtSha1Sum("32c96a03dc8cd20097940f351bca6261ee5a1643"
 // echo -n -e "abcdefgh\n" | sha1sum
 static const std::string kBTxtSha1Sum("e414af7161c9554089f4106d6f1797ef14a73666");
 
-static const char* data_root = getenv("ANDROID_DATA");
-
 static std::string from_testdata_base(const std::string& fname) {
-  return std::string(data_root) + "/nativetest/recovery/testdata/" + fname;
+#ifdef __ANDROID__
+  static std::string data_root = getenv("ANDROID_DATA");
+#else
+  static std::string data_root = std::string(getenv("ANDROID_PRODUCT_OUT")) + "/data";
+#endif
+
+  return data_root + "/nativetest/recovery/testdata/" + fname;
 }
 
 #endif  // _OTA_TEST_CONSTANTS_H
diff --git a/tests/component/applypatch_test.cpp b/tests/component/applypatch_test.cpp
index 016fed9b1..61e06adb6 100644
--- a/tests/component/applypatch_test.cpp
+++ b/tests/component/applypatch_test.cpp
@@ -30,12 +30,16 @@
 #include <android-base/file.h>
 #include <android-base/stringprintf.h>
 #include <android-base/test_utils.h>
+#include <bsdiff/bsdiff.h>
 #include <openssl/sha.h>
 
 #include "applypatch/applypatch.h"
 #include "applypatch/applypatch_modes.h"
 #include "common/test_constants.h"
-#include "print_sha1.h"
+#include "otautil/cache_location.h"
+#include "otautil/print_sha1.h"
+
+using namespace std::string_literals;
 
 static void sha1sum(const std::string& fname, std::string* sha1, size_t* fsize = nullptr) {
   ASSERT_NE(nullptr, sha1);
@@ -53,34 +57,20 @@ static void sha1sum(const std::string& fname, std::string* sha1, size_t* fsize =
 }
 
 static void mangle_file(const std::string& fname) {
-  std::string content;
-  content.reserve(1024);
+  std::string content(1024, '\0');
   for (size_t i = 0; i < 1024; i++) {
     content[i] = rand() % 256;
   }
   ASSERT_TRUE(android::base::WriteStringToFile(content, fname));
 }
 
-static bool file_cmp(const std::string& f1, const std::string& f2) {
-  std::string c1;
-  android::base::ReadFileToString(f1, &c1);
-  std::string c2;
-  android::base::ReadFileToString(f2, &c2);
-  return c1 == c2;
-}
-
 class ApplyPatchTest : public ::testing::Test {
  public:
-  static void SetUpTestCase() {
+  virtual void SetUp() override {
     // set up files
     old_file = from_testdata_base("old.file");
     new_file = from_testdata_base("new.file");
-    patch_file = from_testdata_base("patch.bsdiff");
-    rand_file = "/cache/applypatch_test_rand.file";
-    cache_file = "/cache/saved.file";
-
-    // write stuff to rand_file
-    ASSERT_TRUE(android::base::WriteStringToFile("hello", rand_file));
+    nonexistent_file = from_testdata_base("nonexistent.file");
 
     // set up SHA constants
     sha1sum(old_file, &old_sha1, &old_size);
@@ -90,56 +80,35 @@ class ApplyPatchTest : public ::testing::Test {
     bad_sha1_b = android::base::StringPrintf("%040x", rand());
   }
 
-  static std::string old_file;
-  static std::string new_file;
-  static std::string rand_file;
-  static std::string cache_file;
-  static std::string patch_file;
+  std::string old_file;
+  std::string new_file;
+  std::string nonexistent_file;
 
-  static std::string old_sha1;
-  static std::string new_sha1;
-  static std::string bad_sha1_a;
-  static std::string bad_sha1_b;
+  std::string old_sha1;
+  std::string new_sha1;
+  std::string bad_sha1_a;
+  std::string bad_sha1_b;
 
-  static size_t old_size;
-  static size_t new_size;
+  size_t old_size;
+  size_t new_size;
 };
 
-static void cp(const std::string& src, const std::string& tgt) {
-  std::string cmd = "cp " + src + " " + tgt;
-  system(cmd.c_str());
-}
-
-static void backup_old() {
-  cp(ApplyPatchTest::old_file, ApplyPatchTest::cache_file);
-}
-
-static void restore_old() {
-  cp(ApplyPatchTest::cache_file, ApplyPatchTest::old_file);
-}
-
 class ApplyPatchCacheTest : public ApplyPatchTest {
- public:
-  virtual void SetUp() {
-    backup_old();
+ protected:
+  void SetUp() override {
+    ApplyPatchTest::SetUp();
+    CacheLocation::location().set_cache_temp_source(old_file);
   }
+};
 
-  virtual void TearDown() {
-    restore_old();
+class ApplyPatchModesTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    CacheLocation::location().set_cache_temp_source(cache_source.path);
   }
-};
 
-std::string ApplyPatchTest::old_file;
-std::string ApplyPatchTest::new_file;
-std::string ApplyPatchTest::rand_file;
-std::string ApplyPatchTest::patch_file;
-std::string ApplyPatchTest::cache_file;
-std::string ApplyPatchTest::old_sha1;
-std::string ApplyPatchTest::new_sha1;
-std::string ApplyPatchTest::bad_sha1_a;
-std::string ApplyPatchTest::bad_sha1_b;
-size_t ApplyPatchTest::old_size;
-size_t ApplyPatchTest::new_size;
+  TemporaryFile cache_source;
+};
 
 TEST_F(ApplyPatchTest, CheckModeSkip) {
   std::vector<std::string> sha1s;
@@ -197,43 +166,31 @@ TEST_F(ApplyPatchTest, CheckModeEmmcTarget) {
   ASSERT_EQ(0, applypatch_check(src_file.c_str(), sha1s));
 }
 
-TEST_F(ApplyPatchCacheTest, CheckCacheCorruptedSingle) {
-  mangle_file(old_file);
-  std::vector<std::string> sha1s = { old_sha1 };
-  ASSERT_EQ(0, applypatch_check(&old_file[0], sha1s));
+TEST_F(ApplyPatchCacheTest, CheckCacheCorruptedSourceSingle) {
+  TemporaryFile temp_file;
+  mangle_file(temp_file.path);
+  std::vector<std::string> sha1s_single = { old_sha1 };
+  ASSERT_EQ(0, applypatch_check(temp_file.path, sha1s_single));
+  ASSERT_EQ(0, applypatch_check(nonexistent_file.c_str(), sha1s_single));
 }
 
-TEST_F(ApplyPatchCacheTest, CheckCacheCorruptedMultiple) {
-  mangle_file(old_file);
-  std::vector<std::string> sha1s = { bad_sha1_a, old_sha1, bad_sha1_b };
-  ASSERT_EQ(0, applypatch_check(&old_file[0], sha1s));
+TEST_F(ApplyPatchCacheTest, CheckCacheCorruptedSourceMultiple) {
+  TemporaryFile temp_file;
+  mangle_file(temp_file.path);
+  std::vector<std::string> sha1s_multiple = { bad_sha1_a, old_sha1, bad_sha1_b };
+  ASSERT_EQ(0, applypatch_check(temp_file.path, sha1s_multiple));
+  ASSERT_EQ(0, applypatch_check(nonexistent_file.c_str(), sha1s_multiple));
 }
 
-TEST_F(ApplyPatchCacheTest, CheckCacheCorruptedFailure) {
-  mangle_file(old_file);
-  std::vector<std::string> sha1s = { bad_sha1_a, bad_sha1_b };
-  ASSERT_NE(0, applypatch_check(&old_file[0], sha1s));
+TEST_F(ApplyPatchCacheTest, CheckCacheCorruptedSourceFailure) {
+  TemporaryFile temp_file;
+  mangle_file(temp_file.path);
+  std::vector<std::string> sha1s_failure = { bad_sha1_a, bad_sha1_b };
+  ASSERT_NE(0, applypatch_check(temp_file.path, sha1s_failure));
+  ASSERT_NE(0, applypatch_check(nonexistent_file.c_str(), sha1s_failure));
 }
 
-TEST_F(ApplyPatchCacheTest, CheckCacheMissingSingle) {
-  unlink(&old_file[0]);
-  std::vector<std::string> sha1s = { old_sha1 };
-  ASSERT_EQ(0, applypatch_check(&old_file[0], sha1s));
-}
-
-TEST_F(ApplyPatchCacheTest, CheckCacheMissingMultiple) {
-  unlink(&old_file[0]);
-  std::vector<std::string> sha1s = { bad_sha1_a, old_sha1, bad_sha1_b };
-  ASSERT_EQ(0, applypatch_check(&old_file[0], sha1s));
-}
-
-TEST_F(ApplyPatchCacheTest, CheckCacheMissingFailure) {
-  unlink(&old_file[0]);
-  std::vector<std::string> sha1s = { bad_sha1_a, bad_sha1_b };
-  ASSERT_NE(0, applypatch_check(&old_file[0], sha1s));
-}
-
-TEST(ApplyPatchModesTest, InvalidArgs) {
+TEST_F(ApplyPatchModesTest, InvalidArgs) {
   // At least two args (including the filename).
   ASSERT_EQ(2, applypatch_modes(1, (const char* []){ "applypatch" }));
 
@@ -241,7 +198,7 @@ TEST(ApplyPatchModesTest, InvalidArgs) {
   ASSERT_EQ(2, applypatch_modes(2, (const char* []){ "applypatch", "-x" }));
 }
 
-TEST(ApplyPatchModesTest, PatchModeEmmcTarget) {
+TEST_F(ApplyPatchModesTest, PatchModeEmmcTarget) {
   std::string boot_img = from_testdata_base("boot.img");
   size_t boot_img_size;
   std::string boot_img_sha1;
@@ -311,7 +268,55 @@ TEST(ApplyPatchModesTest, PatchModeEmmcTarget) {
   ASSERT_EQ(0, applypatch_modes(args3.size(), args3.data()));
 }
 
-TEST(ApplyPatchModesTest, PatchModeInvalidArgs) {
+// Ensures that applypatch works with a bsdiff based recovery-from-boot.p.
+TEST_F(ApplyPatchModesTest, PatchModeEmmcTargetWithBsdiffPatch) {
+  std::string boot_img_file = from_testdata_base("boot.img");
+  std::string boot_img_sha1;
+  size_t boot_img_size;
+  sha1sum(boot_img_file, &boot_img_sha1, &boot_img_size);
+
+  std::string recovery_img_file = from_testdata_base("recovery.img");
+  std::string recovery_img_sha1;
+  size_t recovery_img_size;
+  sha1sum(recovery_img_file, &recovery_img_sha1, &recovery_img_size);
+
+  // Generate the bsdiff patch of recovery-from-boot.p.
+  std::string src_content;
+  ASSERT_TRUE(android::base::ReadFileToString(boot_img_file, &src_content));
+
+  std::string tgt_content;
+  ASSERT_TRUE(android::base::ReadFileToString(recovery_img_file, &tgt_content));
+
+  TemporaryFile patch_file;
+  ASSERT_EQ(0,
+            bsdiff::bsdiff(reinterpret_cast<const uint8_t*>(src_content.data()), src_content.size(),
+                           reinterpret_cast<const uint8_t*>(tgt_content.data()), tgt_content.size(),
+                           patch_file.path, nullptr));
+
+  // applypatch <src-file> <tgt-file> <tgt-sha1> <tgt-size> <src-sha1>:<patch>
+  std::string src_file_arg =
+      "EMMC:" + boot_img_file + ":" + std::to_string(boot_img_size) + ":" + boot_img_sha1;
+  TemporaryFile tgt_file;
+  std::string tgt_file_arg = "EMMC:"s + tgt_file.path;
+  std::string recovery_img_size_arg = std::to_string(recovery_img_size);
+  std::string patch_arg = boot_img_sha1 + ":" + patch_file.path;
+  std::vector<const char*> args = { "applypatch",
+                                    src_file_arg.c_str(),
+                                    tgt_file_arg.c_str(),
+                                    recovery_img_sha1.c_str(),
+                                    recovery_img_size_arg.c_str(),
+                                    patch_arg.c_str() };
+  ASSERT_EQ(0, applypatch_modes(args.size(), args.data()));
+
+  // Double check the patched recovery image.
+  std::string tgt_file_sha1;
+  size_t tgt_file_size;
+  sha1sum(tgt_file.path, &tgt_file_sha1, &tgt_file_size);
+  ASSERT_EQ(recovery_img_size, tgt_file_size);
+  ASSERT_EQ(recovery_img_sha1, tgt_file_sha1);
+}
+
+TEST_F(ApplyPatchModesTest, PatchModeInvalidArgs) {
   // Invalid bonus file.
   ASSERT_NE(0, applypatch_modes(3, (const char* []){ "applypatch", "-b", "/doesntexist" }));
 
@@ -372,11 +377,11 @@ TEST(ApplyPatchModesTest, PatchModeInvalidArgs) {
   ASSERT_NE(0, applypatch_modes(args6.size(), args6.data()));
 }
 
-TEST(ApplyPatchModesTest, CheckModeInvalidArgs) {
+TEST_F(ApplyPatchModesTest, CheckModeInvalidArgs) {
   // Insufficient args.
   ASSERT_EQ(2, applypatch_modes(2, (const char* []){ "applypatch", "-c" }));
 }
 
-TEST(ApplyPatchModesTest, ShowLicenses) {
+TEST_F(ApplyPatchModesTest, ShowLicenses) {
   ASSERT_EQ(0, applypatch_modes(2, (const char* []){ "applypatch", "-l" }));
 }
diff --git a/tests/component/bootloader_message_test.cpp b/tests/component/bootloader_message_test.cpp
index b38bc7134..6cc59a495 100644
--- a/tests/component/bootloader_message_test.cpp
+++ b/tests/component/bootloader_message_test.cpp
@@ -18,53 +18,12 @@
 #include <vector>
 
 #include <android-base/strings.h>
+#include <android-base/test_utils.h>
 #include <bootloader_message/bootloader_message.h>
 #include <gtest/gtest.h>
 
-class BootloaderMessageTest : public ::testing::Test {
- protected:
-  BootloaderMessageTest() : has_misc(true) {}
-
-  virtual void SetUp() override {
-    std::string err;
-    has_misc = !get_bootloader_message_blk_device(&err).empty();
-  }
-
-  virtual void TearDown() override {
-    // Clear the BCB.
-    if (has_misc) {
-      std::string err;
-      ASSERT_TRUE(clear_bootloader_message(&err)) << "Failed to clear BCB: " << err;
-    }
-  }
-
-  bool has_misc;
-};
-
-TEST_F(BootloaderMessageTest, clear_bootloader_message) {
-  if (!has_misc) {
-    GTEST_LOG_(INFO) << "Test skipped due to no /misc partition found on the device.";
-    return;
-  }
-
-  // Clear the BCB.
-  std::string err;
-  ASSERT_TRUE(clear_bootloader_message(&err)) << "Failed to clear BCB: " << err;
-
-  // Verify the content.
-  bootloader_message boot;
-  ASSERT_TRUE(read_bootloader_message(&boot, &err)) << "Failed to read BCB: " << err;
-
-  // All the bytes should be cleared.
-  ASSERT_EQ(std::string(sizeof(boot), '\0'),
-            std::string(reinterpret_cast<const char*>(&boot), sizeof(boot)));
-}
-
-TEST_F(BootloaderMessageTest, read_and_write_bootloader_message) {
-  if (!has_misc) {
-    GTEST_LOG_(INFO) << "Test skipped due to no /misc partition found on the device.";
-    return;
-  }
+TEST(BootloaderMessageTest, read_and_write_bootloader_message) {
+  TemporaryFile temp_misc;
 
   // Write the BCB.
   bootloader_message boot = {};
@@ -73,90 +32,71 @@ TEST_F(BootloaderMessageTest, read_and_write_bootloader_message) {
   strlcpy(boot.status, "status1", sizeof(boot.status));
 
   std::string err;
-  ASSERT_TRUE(write_bootloader_message(boot, &err)) << "Failed to write BCB: " << err;
+  ASSERT_TRUE(write_bootloader_message_to(boot, temp_misc.path, &err))
+      << "Failed to write BCB: " << err;
 
   // Read and verify.
   bootloader_message boot_verify;
-  ASSERT_TRUE(read_bootloader_message(&boot_verify, &err)) << "Failed to read BCB: " << err;
+  ASSERT_TRUE(read_bootloader_message_from(&boot_verify, temp_misc.path, &err))
+      << "Failed to read BCB: " << err;
 
   ASSERT_EQ(std::string(reinterpret_cast<const char*>(&boot), sizeof(boot)),
             std::string(reinterpret_cast<const char*>(&boot_verify), sizeof(boot_verify)));
 }
 
-TEST_F(BootloaderMessageTest, write_bootloader_message_options) {
-  if (!has_misc) {
-    GTEST_LOG_(INFO) << "Test skipped due to no /misc partition found on the device.";
-    return;
-  }
-
+TEST(BootloaderMessageTest, update_bootloader_message_in_struct) {
   // Write the options to BCB.
   std::vector<std::string> options = { "option1", "option2" };
-  std::string err;
-  ASSERT_TRUE(write_bootloader_message(options, &err)) << "Failed to write BCB: " << err;
 
-  // Inject some bytes into boot, which should be overwritten while reading.
-  bootloader_message boot;
+  bootloader_message boot = {};
+  // Inject some bytes into boot.
   strlcpy(boot.recovery, "random message", sizeof(boot.recovery));
+  strlcpy(boot.status, "status bytes", sizeof(boot.status));
+  strlcpy(boot.stage, "stage bytes", sizeof(boot.stage));
   strlcpy(boot.reserved, "reserved bytes", sizeof(boot.reserved));
 
-  ASSERT_TRUE(read_bootloader_message(&boot, &err)) << "Failed to read BCB: " << err;
+  ASSERT_TRUE(update_bootloader_message_in_struct(&boot, options));
 
   // Verify that command and recovery fields should be set.
   ASSERT_EQ("boot-recovery", std::string(boot.command));
   std::string expected = "recovery\n" + android::base::Join(options, "\n") + "\n";
   ASSERT_EQ(expected, std::string(boot.recovery));
 
-  // The rest should be cleared.
-  ASSERT_EQ(std::string(sizeof(boot.status), '\0'), std::string(boot.status, sizeof(boot.status)));
-  ASSERT_EQ(std::string(sizeof(boot.stage), '\0'), std::string(boot.stage, sizeof(boot.stage)));
-  ASSERT_EQ(std::string(sizeof(boot.reserved), '\0'),
-            std::string(boot.reserved, sizeof(boot.reserved)));
+  // The rest should be intact.
+  ASSERT_EQ("status bytes", std::string(boot.status));
+  ASSERT_EQ("stage bytes", std::string(boot.stage));
+  ASSERT_EQ("reserved bytes", std::string(boot.reserved));
 }
 
-TEST_F(BootloaderMessageTest, write_bootloader_message_options_empty) {
-  if (!has_misc) {
-    GTEST_LOG_(INFO) << "Test skipped due to no /misc partition found on the device.";
-    return;
-  }
-
+TEST(BootloaderMessageTest, update_bootloader_message_recovery_options_empty) {
   // Write empty vector.
   std::vector<std::string> options;
-  std::string err;
-  ASSERT_TRUE(write_bootloader_message(options, &err)) << "Failed to write BCB: " << err;
 
   // Read and verify.
-  bootloader_message boot;
-  ASSERT_TRUE(read_bootloader_message(&boot, &err)) << "Failed to read BCB: " << err;
+  bootloader_message boot = {};
+  ASSERT_TRUE(update_bootloader_message_in_struct(&boot, options));
 
   // command and recovery fields should be set.
   ASSERT_EQ("boot-recovery", std::string(boot.command));
   ASSERT_EQ("recovery\n", std::string(boot.recovery));
 
-  // The rest should be cleared.
+  // The rest should be empty.
   ASSERT_EQ(std::string(sizeof(boot.status), '\0'), std::string(boot.status, sizeof(boot.status)));
   ASSERT_EQ(std::string(sizeof(boot.stage), '\0'), std::string(boot.stage, sizeof(boot.stage)));
   ASSERT_EQ(std::string(sizeof(boot.reserved), '\0'),
             std::string(boot.reserved, sizeof(boot.reserved)));
 }
 
-TEST_F(BootloaderMessageTest, write_bootloader_message_options_long) {
-  if (!has_misc) {
-    GTEST_LOG_(INFO) << "Test skipped due to no /misc partition found on the device.";
-    return;
-  }
-
+TEST(BootloaderMessageTest, update_bootloader_message_recovery_options_long) {
   // Write super long message.
   std::vector<std::string> options;
   for (int i = 0; i < 100; i++) {
     options.push_back("option: " + std::to_string(i));
   }
 
-  std::string err;
-  ASSERT_TRUE(write_bootloader_message(options, &err)) << "Failed to write BCB: " << err;
-
   // Read and verify.
-  bootloader_message boot;
-  ASSERT_TRUE(read_bootloader_message(&boot, &err)) << "Failed to read BCB: " << err;
+  bootloader_message boot = {};
+  ASSERT_TRUE(update_bootloader_message_in_struct(&boot, options));
 
   // Make sure it's long enough.
   std::string expected = "recovery\n" + android::base::Join(options, "\n") + "\n";
@@ -167,40 +107,10 @@ TEST_F(BootloaderMessageTest, write_bootloader_message_options_long) {
   ASSERT_EQ(expected.substr(0, sizeof(boot.recovery) - 1), std::string(boot.recovery));
   ASSERT_EQ('\0', boot.recovery[sizeof(boot.recovery) - 1]);
 
-  // The rest should be cleared.
+  // The rest should be empty.
   ASSERT_EQ(std::string(sizeof(boot.status), '\0'), std::string(boot.status, sizeof(boot.status)));
   ASSERT_EQ(std::string(sizeof(boot.stage), '\0'), std::string(boot.stage, sizeof(boot.stage)));
   ASSERT_EQ(std::string(sizeof(boot.reserved), '\0'),
             std::string(boot.reserved, sizeof(boot.reserved)));
 }
 
-TEST_F(BootloaderMessageTest, update_bootloader_message) {
-  if (!has_misc) {
-    GTEST_LOG_(INFO) << "Test skipped due to no /misc partition found on the device.";
-    return;
-  }
-
-  // Inject some bytes into boot, which should be not overwritten later.
-  bootloader_message boot;
-  strlcpy(boot.recovery, "random message", sizeof(boot.recovery));
-  strlcpy(boot.reserved, "reserved bytes", sizeof(boot.reserved));
-  std::string err;
-  ASSERT_TRUE(write_bootloader_message(boot, &err)) << "Failed to write BCB: " << err;
-
-  // Update the BCB message.
-  std::vector<std::string> options = { "option1", "option2" };
-  ASSERT_TRUE(update_bootloader_message(options, &err)) << "Failed to update BCB: " << err;
-
-  bootloader_message boot_verify;
-  ASSERT_TRUE(read_bootloader_message(&boot_verify, &err)) << "Failed to read BCB: " << err;
-
-  // Verify that command and recovery fields should be set.
-  ASSERT_EQ("boot-recovery", std::string(boot_verify.command));
-  std::string expected = "recovery\n" + android::base::Join(options, "\n") + "\n";
-  ASSERT_EQ(expected, std::string(boot_verify.recovery));
-
-  // The rest should be intact.
-  ASSERT_EQ(std::string(boot.status), std::string(boot_verify.status));
-  ASSERT_EQ(std::string(boot.stage), std::string(boot_verify.stage));
-  ASSERT_EQ(std::string(boot.reserved), std::string(boot_verify.reserved));
-}
diff --git a/tests/component/imgdiff_test.cpp b/tests/component/imgdiff_test.cpp
index bf25aebb0..6c23def01 100644
--- a/tests/component/imgdiff_test.cpp
+++ b/tests/component/imgdiff_test.cpp
@@ -16,17 +16,24 @@
 
 #include <stdio.h>
 
+#include <algorithm>
 #include <string>
+#include <tuple>
 #include <vector>
 
 #include <android-base/file.h>
 #include <android-base/memory.h>
+#include <android-base/stringprintf.h>
+#include <android-base/strings.h>
 #include <android-base/test_utils.h>
 #include <applypatch/imgdiff.h>
+#include <applypatch/imgdiff_image.h>
 #include <applypatch/imgpatch.h>
 #include <gtest/gtest.h>
 #include <ziparchive/zip_writer.h>
 
+#include "common/test_constants.h"
+
 using android::base::get_unaligned;
 
 // Sanity check for the given imgdiff patch header.
@@ -75,15 +82,20 @@ static void verify_patch_header(const std::string& patch, size_t* num_normal, si
   if (num_deflate != nullptr) *num_deflate = deflate;
 }
 
-static void verify_patched_image(const std::string& src, const std::string& patch,
-                                 const std::string& tgt) {
-  std::string patched;
+static void GenerateTarget(const std::string& src, const std::string& patch, std::string* patched) {
+  patched->clear();
   ASSERT_EQ(0, ApplyImagePatch(reinterpret_cast<const unsigned char*>(src.data()), src.size(),
                                reinterpret_cast<const unsigned char*>(patch.data()), patch.size(),
-                               [&patched](const unsigned char* data, size_t len) {
-                                 patched.append(reinterpret_cast<const char*>(data), len);
+                               [&](const unsigned char* data, size_t len) {
+                                 patched->append(reinterpret_cast<const char*>(data), len);
                                  return len;
                                }));
+}
+
+static void verify_patched_image(const std::string& src, const std::string& patch,
+                                 const std::string& tgt) {
+  std::string patched;
+  GenerateTarget(src, patch, &patched);
   ASSERT_EQ(tgt, patched);
 }
 
@@ -138,7 +150,7 @@ TEST(ImgdiffTest, image_mode_smoke) {
 TEST(ImgdiffTest, zip_mode_smoke_store) {
   // Construct src and tgt zip files.
   TemporaryFile src_file;
-  FILE* src_file_ptr = fdopen(src_file.fd, "wb");
+  FILE* src_file_ptr = fdopen(src_file.release(), "wb");
   ZipWriter src_writer(src_file_ptr);
   ASSERT_EQ(0, src_writer.StartEntry("file1.txt", 0));  // Store mode.
   const std::string src_content("abcdefg");
@@ -148,7 +160,7 @@ TEST(ImgdiffTest, zip_mode_smoke_store) {
   ASSERT_EQ(0, fclose(src_file_ptr));
 
   TemporaryFile tgt_file;
-  FILE* tgt_file_ptr = fdopen(tgt_file.fd, "wb");
+  FILE* tgt_file_ptr = fdopen(tgt_file.release(), "wb");
   ZipWriter tgt_writer(tgt_file_ptr);
   ASSERT_EQ(0, tgt_writer.StartEntry("file1.txt", 0));  // Store mode.
   const std::string tgt_content("abcdefgxyz");
@@ -187,7 +199,7 @@ TEST(ImgdiffTest, zip_mode_smoke_store) {
 TEST(ImgdiffTest, zip_mode_smoke_compressed) {
   // Construct src and tgt zip files.
   TemporaryFile src_file;
-  FILE* src_file_ptr = fdopen(src_file.fd, "wb");
+  FILE* src_file_ptr = fdopen(src_file.release(), "wb");
   ZipWriter src_writer(src_file_ptr);
   ASSERT_EQ(0, src_writer.StartEntry("file1.txt", ZipWriter::kCompress));
   const std::string src_content("abcdefg");
@@ -197,7 +209,7 @@ TEST(ImgdiffTest, zip_mode_smoke_compressed) {
   ASSERT_EQ(0, fclose(src_file_ptr));
 
   TemporaryFile tgt_file;
-  FILE* tgt_file_ptr = fdopen(tgt_file.fd, "wb");
+  FILE* tgt_file_ptr = fdopen(tgt_file.release(), "wb");
   ZipWriter tgt_writer(tgt_file_ptr);
   ASSERT_EQ(0, tgt_writer.StartEntry("file1.txt", ZipWriter::kCompress));
   const std::string tgt_content("abcdefgxyz");
@@ -236,7 +248,7 @@ TEST(ImgdiffTest, zip_mode_smoke_compressed) {
 TEST(ImgdiffTest, zip_mode_smoke_trailer_zeros) {
   // Construct src and tgt zip files.
   TemporaryFile src_file;
-  FILE* src_file_ptr = fdopen(src_file.fd, "wb");
+  FILE* src_file_ptr = fdopen(src_file.release(), "wb");
   ZipWriter src_writer(src_file_ptr);
   ASSERT_EQ(0, src_writer.StartEntry("file1.txt", ZipWriter::kCompress));
   const std::string src_content("abcdefg");
@@ -246,7 +258,7 @@ TEST(ImgdiffTest, zip_mode_smoke_trailer_zeros) {
   ASSERT_EQ(0, fclose(src_file_ptr));
 
   TemporaryFile tgt_file;
-  FILE* tgt_file_ptr = fdopen(tgt_file.fd, "wb");
+  FILE* tgt_file_ptr = fdopen(tgt_file.release(), "wb");
   ZipWriter tgt_writer(tgt_file_ptr);
   ASSERT_EQ(0, tgt_writer.StartEntry("file1.txt", ZipWriter::kCompress));
   const std::string tgt_content("abcdefgxyz");
@@ -623,3 +635,438 @@ TEST(ImgpatchTest, image_mode_patch_corruption) {
                                 reinterpret_cast<const unsigned char*>(patch.data()), patch.size(),
                                 [](const unsigned char* /*data*/, size_t len) { return len; }));
 }
+
+static void construct_store_entry(const std::vector<std::tuple<std::string, size_t, char>>& info,
+                                  ZipWriter* writer) {
+  for (auto& t : info) {
+    // Create t(1) blocks of t(2), and write the data to t(0)
+    ASSERT_EQ(0, writer->StartEntry(std::get<0>(t).c_str(), 0));
+    const std::string content(std::get<1>(t) * 4096, std::get<2>(t));
+    ASSERT_EQ(0, writer->WriteBytes(content.data(), content.size()));
+    ASSERT_EQ(0, writer->FinishEntry());
+  }
+}
+
+static void construct_deflate_entry(const std::vector<std::tuple<std::string, size_t, size_t>>& info,
+                                    ZipWriter* writer, const std::string& data) {
+  for (auto& t : info) {
+    // t(0): entry_name; t(1): block offset; t(2) length in blocks.
+    ASSERT_EQ(0, writer->StartEntry(std::get<0>(t).c_str(), ZipWriter::kCompress));
+    ASSERT_EQ(0, writer->WriteBytes(data.data() + std::get<1>(t) * 4096, std::get<2>(t) * 4096));
+    ASSERT_EQ(0, writer->FinishEntry());
+  }
+}
+
+// Look for the source and patch pieces in debug_dir. Generate a target piece from each pair.
+// Concatenate all the target pieces and match against the orignal one. Used pieces in debug_dir
+// will be cleaned up.
+static void GenerateAndCheckSplitTarget(const std::string& debug_dir, size_t count,
+                                        const std::string& tgt) {
+  std::string patched;
+  for (size_t i = 0; i < count; i++) {
+    std::string split_src_path = android::base::StringPrintf("%s/src-%zu", debug_dir.c_str(), i);
+    std::string split_src;
+    ASSERT_TRUE(android::base::ReadFileToString(split_src_path, &split_src));
+    ASSERT_EQ(0, unlink(split_src_path.c_str()));
+
+    std::string split_patch_path =
+        android::base::StringPrintf("%s/patch-%zu", debug_dir.c_str(), i);
+    std::string split_patch;
+    ASSERT_TRUE(android::base::ReadFileToString(split_patch_path, &split_patch));
+    ASSERT_EQ(0, unlink(split_patch_path.c_str()));
+
+    std::string split_tgt;
+    GenerateTarget(split_src, split_patch, &split_tgt);
+    patched += split_tgt;
+  }
+
+  // Verify we can get back the original target image.
+  ASSERT_EQ(tgt, patched);
+}
+
+std::vector<ImageChunk> ConstructImageChunks(
+    const std::vector<uint8_t>& content, const std::vector<std::tuple<std::string, size_t>>& info) {
+  std::vector<ImageChunk> chunks;
+  size_t start = 0;
+  for (const auto& t : info) {
+    size_t length = std::get<1>(t);
+    chunks.emplace_back(CHUNK_NORMAL, start, &content, length, std::get<0>(t));
+    start += length;
+  }
+
+  return chunks;
+}
+
+TEST(ImgdiffTest, zip_mode_split_image_smoke) {
+  std::vector<uint8_t> content;
+  content.reserve(4096 * 50);
+  uint8_t n = 0;
+  generate_n(back_inserter(content), 4096 * 50, [&n]() { return n++ / 4096; });
+
+  ZipModeImage tgt_image(false, 4096 * 10);
+  std::vector<ImageChunk> tgt_chunks = ConstructImageChunks(content, { { "a", 100 },
+                                                                       { "b", 4096 * 2 },
+                                                                       { "c", 4096 * 3 },
+                                                                       { "d", 300 },
+                                                                       { "e-0", 4096 * 10 },
+                                                                       { "e-1", 4096 * 5 },
+                                                                       { "CD", 200 } });
+  tgt_image.Initialize(std::move(tgt_chunks),
+                       std::vector<uint8_t>(content.begin(), content.begin() + 82520));
+
+  tgt_image.DumpChunks();
+
+  ZipModeImage src_image(true, 4096 * 10);
+  std::vector<ImageChunk> src_chunks = ConstructImageChunks(content, { { "b", 4096 * 3 },
+                                                                       { "c-0", 4096 * 10 },
+                                                                       { "c-1", 4096 * 2 },
+                                                                       { "a", 4096 * 5 },
+                                                                       { "e-0", 4096 * 10 },
+                                                                       { "e-1", 10000 },
+                                                                       { "CD", 5000 } });
+  src_image.Initialize(std::move(src_chunks),
+                       std::vector<uint8_t>(content.begin(), content.begin() + 137880));
+
+  std::vector<ZipModeImage> split_tgt_images;
+  std::vector<ZipModeImage> split_src_images;
+  std::vector<SortedRangeSet> split_src_ranges;
+
+  ZipModeImage::SplitZipModeImageWithLimit(tgt_image, src_image, &split_tgt_images,
+                                           &split_src_images, &split_src_ranges);
+
+  // src_piece 1: a 5 blocks, b 3 blocks
+  // src_piece 2: c-0 10 blocks
+  // src_piece 3: d 0 block, e-0 10 blocks
+  // src_piece 4: e-1 2 blocks; CD 2 blocks
+  ASSERT_EQ(split_tgt_images.size(), split_src_images.size());
+  ASSERT_EQ(static_cast<size_t>(4), split_tgt_images.size());
+
+  ASSERT_EQ(static_cast<size_t>(1), split_tgt_images[0].NumOfChunks());
+  ASSERT_EQ(static_cast<size_t>(12288), split_tgt_images[0][0].DataLengthForPatch());
+  ASSERT_EQ("4,0,3,15,20", split_src_ranges[0].ToString());
+
+  ASSERT_EQ(static_cast<size_t>(1), split_tgt_images[1].NumOfChunks());
+  ASSERT_EQ(static_cast<size_t>(12288), split_tgt_images[1][0].DataLengthForPatch());
+  ASSERT_EQ("2,3,13", split_src_ranges[1].ToString());
+
+  ASSERT_EQ(static_cast<size_t>(1), split_tgt_images[2].NumOfChunks());
+  ASSERT_EQ(static_cast<size_t>(40960), split_tgt_images[2][0].DataLengthForPatch());
+  ASSERT_EQ("2,20,30", split_src_ranges[2].ToString());
+
+  ASSERT_EQ(static_cast<size_t>(1), split_tgt_images[3].NumOfChunks());
+  ASSERT_EQ(static_cast<size_t>(16984), split_tgt_images[3][0].DataLengthForPatch());
+  ASSERT_EQ("2,30,34", split_src_ranges[3].ToString());
+}
+
+TEST(ImgdiffTest, zip_mode_store_large_apk) {
+  // Construct src and tgt zip files with limit = 10 blocks.
+  //     src              tgt
+  //  12 blocks 'd'    3 blocks  'a'
+  //  8 blocks  'c'    3 blocks  'b'
+  //  3 blocks  'b'    8 blocks  'c' (exceeds limit)
+  //  3 blocks  'a'    12 blocks 'd' (exceeds limit)
+  //                   3 blocks  'e'
+  TemporaryFile tgt_file;
+  FILE* tgt_file_ptr = fdopen(tgt_file.release(), "wb");
+  ZipWriter tgt_writer(tgt_file_ptr);
+  construct_store_entry(
+      { { "a", 3, 'a' }, { "b", 3, 'b' }, { "c", 8, 'c' }, { "d", 12, 'd' }, { "e", 3, 'e' } },
+      &tgt_writer);
+  ASSERT_EQ(0, tgt_writer.Finish());
+  ASSERT_EQ(0, fclose(tgt_file_ptr));
+
+  TemporaryFile src_file;
+  FILE* src_file_ptr = fdopen(src_file.release(), "wb");
+  ZipWriter src_writer(src_file_ptr);
+  construct_store_entry({ { "d", 12, 'd' }, { "c", 8, 'c' }, { "b", 3, 'b' }, { "a", 3, 'a' } },
+                        &src_writer);
+  ASSERT_EQ(0, src_writer.Finish());
+  ASSERT_EQ(0, fclose(src_file_ptr));
+
+  // Compute patch.
+  TemporaryFile patch_file;
+  TemporaryFile split_info_file;
+  TemporaryDir debug_dir;
+  std::string split_info_arg = android::base::StringPrintf("--split-info=%s", split_info_file.path);
+  std::string debug_dir_arg = android::base::StringPrintf("--debug-dir=%s", debug_dir.path);
+  std::vector<const char*> args = {
+    "imgdiff", "-z", "--block-limit=10", split_info_arg.c_str(), debug_dir_arg.c_str(),
+    src_file.path, tgt_file.path, patch_file.path,
+  };
+  ASSERT_EQ(0, imgdiff(args.size(), args.data()));
+
+  std::string tgt;
+  ASSERT_TRUE(android::base::ReadFileToString(tgt_file.path, &tgt));
+
+  // Expect 4 pieces of patch. (Roughly 3'a',3'b'; 8'c'; 10'd'; 2'd'3'e')
+  GenerateAndCheckSplitTarget(debug_dir.path, 4, tgt);
+}
+
+TEST(ImgdiffTest, zip_mode_deflate_large_apk) {
+  // Src and tgt zip files are constructed as follows.
+  //     src               tgt
+  //  22 blocks, "d"    4  blocks,  "a"
+  //  5 blocks,  "b"    4  blocks,  "b"
+  //  3 blocks,  "a"    8  blocks,  "c" (exceeds limit)
+  //  1 block,   "g"    20 blocks,  "d" (exceeds limit)
+  //  8 blocks,  "c"    2  blocks,  "e"
+  //  1 block,   "f"    1  block ,  "f"
+  std::string tgt_path = from_testdata_base("deflate_tgt.zip");
+  std::string src_path = from_testdata_base("deflate_src.zip");
+
+  ZipModeImage src_image(true, 10 * 4096);
+  ZipModeImage tgt_image(false, 10 * 4096);
+  ASSERT_TRUE(src_image.Initialize(src_path));
+  ASSERT_TRUE(tgt_image.Initialize(tgt_path));
+  ASSERT_TRUE(ZipModeImage::CheckAndProcessChunks(&tgt_image, &src_image));
+
+  src_image.DumpChunks();
+  tgt_image.DumpChunks();
+
+  std::vector<ZipModeImage> split_tgt_images;
+  std::vector<ZipModeImage> split_src_images;
+  std::vector<SortedRangeSet> split_src_ranges;
+  ZipModeImage::SplitZipModeImageWithLimit(tgt_image, src_image, &split_tgt_images,
+                                           &split_src_images, &split_src_ranges);
+
+  // Expected split images with limit = 10 blocks.
+  // src_piece 0: a 3 blocks, b 5 blocks
+  // src_piece 1: c 8 blocks
+  // src_piece 2: d-0 10 block
+  // src_piece 3: d-1 10 blocks
+  // src_piece 4: e 1 block, CD
+  ASSERT_EQ(split_tgt_images.size(), split_src_images.size());
+  ASSERT_EQ(static_cast<size_t>(5), split_tgt_images.size());
+
+  ASSERT_EQ(static_cast<size_t>(2), split_src_images[0].NumOfChunks());
+  ASSERT_EQ("a", split_src_images[0][0].GetEntryName());
+  ASSERT_EQ("b", split_src_images[0][1].GetEntryName());
+
+  ASSERT_EQ(static_cast<size_t>(1), split_src_images[1].NumOfChunks());
+  ASSERT_EQ("c", split_src_images[1][0].GetEntryName());
+
+  ASSERT_EQ(static_cast<size_t>(0), split_src_images[2].NumOfChunks());
+  ASSERT_EQ(static_cast<size_t>(0), split_src_images[3].NumOfChunks());
+  ASSERT_EQ(static_cast<size_t>(0), split_src_images[4].NumOfChunks());
+
+  // Compute patch.
+  TemporaryFile patch_file;
+  TemporaryFile split_info_file;
+  TemporaryDir debug_dir;
+  ASSERT_TRUE(ZipModeImage::GeneratePatches(split_tgt_images, split_src_images, split_src_ranges,
+                                            patch_file.path, split_info_file.path, debug_dir.path));
+
+  // Verify the content of split info.
+  // Expect 5 pieces of patch. ["a","b"; "c"; "d-0"; "d-1"; "e"]
+  std::string split_info_string;
+  android::base::ReadFileToString(split_info_file.path, &split_info_string);
+  std::vector<std::string> info_list =
+      android::base::Split(android::base::Trim(split_info_string), "\n");
+
+  ASSERT_EQ(static_cast<size_t>(7), info_list.size());
+  ASSERT_EQ("2", android::base::Trim(info_list[0]));
+  ASSERT_EQ("5", android::base::Trim(info_list[1]));
+
+  std::string tgt;
+  ASSERT_TRUE(android::base::ReadFileToString(tgt_path, &tgt));
+  ASSERT_EQ(static_cast<size_t>(160385), tgt.size());
+  std::vector<std::string> tgt_file_ranges = {
+    "36864 2,22,31", "32768 2,31,40", "40960 2,0,11", "40960 2,11,21", "8833 4,21,22,40,41",
+  };
+
+  for (size_t i = 0; i < 5; i++) {
+    struct stat st;
+    std::string path = android::base::StringPrintf("%s/patch-%zu", debug_dir.path, i);
+    ASSERT_EQ(0, stat(path.c_str(), &st));
+    ASSERT_EQ(std::to_string(st.st_size) + " " + tgt_file_ranges[i],
+              android::base::Trim(info_list[i + 2]));
+  }
+
+  GenerateAndCheckSplitTarget(debug_dir.path, 5, tgt);
+}
+
+TEST(ImgdiffTest, zip_mode_no_match_source) {
+  // Generate 20 blocks of random data.
+  std::string random_data;
+  random_data.reserve(4096 * 20);
+  generate_n(back_inserter(random_data), 4096 * 20, []() { return rand() % 256; });
+
+  TemporaryFile tgt_file;
+  FILE* tgt_file_ptr = fdopen(tgt_file.release(), "wb");
+  ZipWriter tgt_writer(tgt_file_ptr);
+
+  construct_deflate_entry({ { "a", 0, 4 }, { "b", 5, 5 }, { "c", 11, 5 } }, &tgt_writer,
+                          random_data);
+
+  ASSERT_EQ(0, tgt_writer.Finish());
+  ASSERT_EQ(0, fclose(tgt_file_ptr));
+
+  // We don't have a matching source entry.
+  TemporaryFile src_file;
+  FILE* src_file_ptr = fdopen(src_file.release(), "wb");
+  ZipWriter src_writer(src_file_ptr);
+  construct_store_entry({ { "d", 1, 'd' } }, &src_writer);
+  ASSERT_EQ(0, src_writer.Finish());
+  ASSERT_EQ(0, fclose(src_file_ptr));
+
+  // Compute patch.
+  TemporaryFile patch_file;
+  TemporaryFile split_info_file;
+  TemporaryDir debug_dir;
+  std::string split_info_arg = android::base::StringPrintf("--split-info=%s", split_info_file.path);
+  std::string debug_dir_arg = android::base::StringPrintf("--debug-dir=%s", debug_dir.path);
+  std::vector<const char*> args = {
+    "imgdiff", "-z", "--block-limit=10", debug_dir_arg.c_str(), split_info_arg.c_str(),
+    src_file.path, tgt_file.path, patch_file.path,
+  };
+  ASSERT_EQ(0, imgdiff(args.size(), args.data()));
+
+  std::string tgt;
+  ASSERT_TRUE(android::base::ReadFileToString(tgt_file.path, &tgt));
+
+  // Expect 1 pieces of patch due to no matching source entry.
+  GenerateAndCheckSplitTarget(debug_dir.path, 1, tgt);
+}
+
+TEST(ImgdiffTest, zip_mode_large_enough_limit) {
+  // Generate 20 blocks of random data.
+  std::string random_data;
+  random_data.reserve(4096 * 20);
+  generate_n(back_inserter(random_data), 4096 * 20, []() { return rand() % 256; });
+
+  TemporaryFile tgt_file;
+  FILE* tgt_file_ptr = fdopen(tgt_file.release(), "wb");
+  ZipWriter tgt_writer(tgt_file_ptr);
+
+  construct_deflate_entry({ { "a", 0, 10 }, { "b", 10, 5 } }, &tgt_writer, random_data);
+
+  ASSERT_EQ(0, tgt_writer.Finish());
+  ASSERT_EQ(0, fclose(tgt_file_ptr));
+
+  // Construct 10 blocks of source.
+  TemporaryFile src_file;
+  FILE* src_file_ptr = fdopen(src_file.release(), "wb");
+  ZipWriter src_writer(src_file_ptr);
+  construct_deflate_entry({ { "a", 1, 10 } }, &src_writer, random_data);
+  ASSERT_EQ(0, src_writer.Finish());
+  ASSERT_EQ(0, fclose(src_file_ptr));
+
+  // Compute patch with a limit of 20 blocks.
+  TemporaryFile patch_file;
+  TemporaryFile split_info_file;
+  TemporaryDir debug_dir;
+  std::string split_info_arg = android::base::StringPrintf("--split-info=%s", split_info_file.path);
+  std::string debug_dir_arg = android::base::StringPrintf("--debug-dir=%s", debug_dir.path);
+  std::vector<const char*> args = {
+    "imgdiff", "-z", "--block-limit=20", split_info_arg.c_str(), debug_dir_arg.c_str(),
+    src_file.path, tgt_file.path, patch_file.path,
+  };
+  ASSERT_EQ(0, imgdiff(args.size(), args.data()));
+
+  std::string tgt;
+  ASSERT_TRUE(android::base::ReadFileToString(tgt_file.path, &tgt));
+
+  // Expect 1 piece of patch since limit is larger than the zip file size.
+  GenerateAndCheckSplitTarget(debug_dir.path, 1, tgt);
+}
+
+TEST(ImgdiffTest, zip_mode_large_apk_small_target_chunk) {
+  TemporaryFile tgt_file;
+  FILE* tgt_file_ptr = fdopen(tgt_file.release(), "wb");
+  ZipWriter tgt_writer(tgt_file_ptr);
+
+  // The first entry is less than 4096 bytes, followed immediately by an entry that has a very
+  // large counterpart in the source file. Therefore the first entry will be patched separately.
+  std::string small_chunk("a", 2000);
+  ASSERT_EQ(0, tgt_writer.StartEntry("a", 0));
+  ASSERT_EQ(0, tgt_writer.WriteBytes(small_chunk.data(), small_chunk.size()));
+  ASSERT_EQ(0, tgt_writer.FinishEntry());
+  construct_store_entry(
+      {
+          { "b", 12, 'b' }, { "c", 3, 'c' },
+      },
+      &tgt_writer);
+  ASSERT_EQ(0, tgt_writer.Finish());
+  ASSERT_EQ(0, fclose(tgt_file_ptr));
+
+  TemporaryFile src_file;
+  FILE* src_file_ptr = fdopen(src_file.release(), "wb");
+  ZipWriter src_writer(src_file_ptr);
+  construct_store_entry({ { "a", 1, 'a' }, { "b", 13, 'b' }, { "c", 1, 'c' } }, &src_writer);
+  ASSERT_EQ(0, src_writer.Finish());
+  ASSERT_EQ(0, fclose(src_file_ptr));
+
+  // Compute patch.
+  TemporaryFile patch_file;
+  TemporaryFile split_info_file;
+  TemporaryDir debug_dir;
+  std::string split_info_arg = android::base::StringPrintf("--split-info=%s", split_info_file.path);
+  std::string debug_dir_arg = android::base::StringPrintf("--debug-dir=%s", debug_dir.path);
+  std::vector<const char*> args = {
+    "imgdiff",     "-z",          "--block-limit=10", split_info_arg.c_str(), debug_dir_arg.c_str(),
+    src_file.path, tgt_file.path, patch_file.path,
+  };
+  ASSERT_EQ(0, imgdiff(args.size(), args.data()));
+
+  std::string tgt;
+  ASSERT_TRUE(android::base::ReadFileToString(tgt_file.path, &tgt));
+
+  // Expect three split src images:
+  // src_piece 0: a 1 blocks
+  // src_piece 1: b-0 10 blocks
+  // src_piece 2: b-1 3 blocks, c 1 blocks, CD
+  GenerateAndCheckSplitTarget(debug_dir.path, 3, tgt);
+}
+
+TEST(ImgdiffTest, zip_mode_large_apk_skipped_small_target_chunk) {
+  TemporaryFile tgt_file;
+  FILE* tgt_file_ptr = fdopen(tgt_file.release(), "wb");
+  ZipWriter tgt_writer(tgt_file_ptr);
+
+  construct_store_entry(
+      {
+          { "a", 11, 'a' },
+      },
+      &tgt_writer);
+
+  // Construct a tiny target entry of 1 byte, which will be skipped due to the tail alignment of
+  // the previous entry.
+  std::string small_chunk("b", 1);
+  ASSERT_EQ(0, tgt_writer.StartEntry("b", 0));
+  ASSERT_EQ(0, tgt_writer.WriteBytes(small_chunk.data(), small_chunk.size()));
+  ASSERT_EQ(0, tgt_writer.FinishEntry());
+
+  ASSERT_EQ(0, tgt_writer.Finish());
+  ASSERT_EQ(0, fclose(tgt_file_ptr));
+
+  TemporaryFile src_file;
+  FILE* src_file_ptr = fdopen(src_file.release(), "wb");
+  ZipWriter src_writer(src_file_ptr);
+  construct_store_entry(
+      {
+          { "a", 11, 'a' }, { "b", 11, 'b' },
+      },
+      &src_writer);
+  ASSERT_EQ(0, src_writer.Finish());
+  ASSERT_EQ(0, fclose(src_file_ptr));
+
+  // Compute patch.
+  TemporaryFile patch_file;
+  TemporaryFile split_info_file;
+  TemporaryDir debug_dir;
+  std::string split_info_arg = android::base::StringPrintf("--split-info=%s", split_info_file.path);
+  std::string debug_dir_arg = android::base::StringPrintf("--debug-dir=%s", debug_dir.path);
+  std::vector<const char*> args = {
+    "imgdiff",     "-z",          "--block-limit=10", split_info_arg.c_str(), debug_dir_arg.c_str(),
+    src_file.path, tgt_file.path, patch_file.path,
+  };
+  ASSERT_EQ(0, imgdiff(args.size(), args.data()));
+
+  std::string tgt;
+  ASSERT_TRUE(android::base::ReadFileToString(tgt_file.path, &tgt));
+
+  // Expect two split src images:
+  // src_piece 0: a-0 10 blocks
+  // src_piece 1: a-0 1 block, CD
+  GenerateAndCheckSplitTarget(debug_dir.path, 2, tgt);
+}
diff --git a/tests/component/install_test.cpp b/tests/component/install_test.cpp
index 968196fc0..d19d788e4 100644
--- a/tests/component/install_test.cpp
+++ b/tests/component/install_test.cpp
@@ -19,6 +19,7 @@
 #include <sys/types.h>
 #include <unistd.h>
 
+#include <algorithm>
 #include <string>
 #include <vector>
 
@@ -36,7 +37,7 @@
 
 TEST(InstallTest, verify_package_compatibility_no_entry) {
   TemporaryFile temp_file;
-  FILE* zip_file = fdopen(temp_file.fd, "w");
+  FILE* zip_file = fdopen(temp_file.release(), "w");
   ZipWriter writer(zip_file);
   // The archive must have something to be opened correctly.
   ASSERT_EQ(0, writer.StartEntry("dummy_entry", 0));
@@ -53,7 +54,7 @@ TEST(InstallTest, verify_package_compatibility_no_entry) {
 
 TEST(InstallTest, verify_package_compatibility_invalid_entry) {
   TemporaryFile temp_file;
-  FILE* zip_file = fdopen(temp_file.fd, "w");
+  FILE* zip_file = fdopen(temp_file.release(), "w");
   ZipWriter writer(zip_file);
   ASSERT_EQ(0, writer.StartEntry("compatibility.zip", 0));
   ASSERT_EQ(0, writer.FinishEntry());
@@ -69,7 +70,7 @@ TEST(InstallTest, verify_package_compatibility_invalid_entry) {
 
 TEST(InstallTest, read_metadata_from_package_smoke) {
   TemporaryFile temp_file;
-  FILE* zip_file = fdopen(temp_file.fd, "w");
+  FILE* zip_file = fdopen(temp_file.release(), "w");
   ZipWriter writer(zip_file);
   ASSERT_EQ(0, writer.StartEntry("META-INF/com/android/metadata", kCompressStored));
   const std::string content("abcdefg");
@@ -86,7 +87,7 @@ TEST(InstallTest, read_metadata_from_package_smoke) {
   CloseArchive(zip);
 
   TemporaryFile temp_file2;
-  FILE* zip_file2 = fdopen(temp_file2.fd, "w");
+  FILE* zip_file2 = fdopen(temp_file2.release(), "w");
   ZipWriter writer2(zip_file2);
   ASSERT_EQ(0, writer2.StartEntry("META-INF/com/android/metadata", kCompressDeflated));
   ASSERT_EQ(0, writer2.WriteBytes(content.data(), content.size()));
@@ -103,7 +104,7 @@ TEST(InstallTest, read_metadata_from_package_smoke) {
 
 TEST(InstallTest, read_metadata_from_package_no_entry) {
   TemporaryFile temp_file;
-  FILE* zip_file = fdopen(temp_file.fd, "w");
+  FILE* zip_file = fdopen(temp_file.release(), "w");
   ZipWriter writer(zip_file);
   ASSERT_EQ(0, writer.StartEntry("dummy_entry", kCompressStored));
   ASSERT_EQ(0, writer.FinishEntry());
@@ -119,7 +120,7 @@ TEST(InstallTest, read_metadata_from_package_no_entry) {
 
 TEST(InstallTest, verify_package_compatibility_with_libvintf_malformed_xml) {
   TemporaryFile compatibility_zip_file;
-  FILE* compatibility_zip = fdopen(compatibility_zip_file.fd, "w");
+  FILE* compatibility_zip = fdopen(compatibility_zip_file.release(), "w");
   ZipWriter compatibility_zip_writer(compatibility_zip);
   ASSERT_EQ(0, compatibility_zip_writer.StartEntry("system_manifest.xml", kCompressDeflated));
   std::string malformed_xml = "malformed";
@@ -129,7 +130,7 @@ TEST(InstallTest, verify_package_compatibility_with_libvintf_malformed_xml) {
   ASSERT_EQ(0, fclose(compatibility_zip));
 
   TemporaryFile temp_file;
-  FILE* zip_file = fdopen(temp_file.fd, "w");
+  FILE* zip_file = fdopen(temp_file.release(), "w");
   ZipWriter writer(zip_file);
   ASSERT_EQ(0, writer.StartEntry("compatibility.zip", kCompressStored));
   std::string compatibility_zip_content;
@@ -164,7 +165,7 @@ TEST(InstallTest, verify_package_compatibility_with_libvintf_system_manifest_xml
   ASSERT_TRUE(
       android::base::ReadFileToString(system_manifest_xml_path, &system_manifest_xml_content));
   TemporaryFile compatibility_zip_file;
-  FILE* compatibility_zip = fdopen(compatibility_zip_file.fd, "w");
+  FILE* compatibility_zip = fdopen(compatibility_zip_file.release(), "w");
   ZipWriter compatibility_zip_writer(compatibility_zip);
   ASSERT_EQ(0, compatibility_zip_writer.StartEntry("system_manifest.xml", kCompressDeflated));
   ASSERT_EQ(0, compatibility_zip_writer.WriteBytes(system_manifest_xml_content.data(),
@@ -174,7 +175,7 @@ TEST(InstallTest, verify_package_compatibility_with_libvintf_system_manifest_xml
   ASSERT_EQ(0, fclose(compatibility_zip));
 
   TemporaryFile temp_file;
-  FILE* zip_file = fdopen(temp_file.fd, "w");
+  FILE* zip_file = fdopen(temp_file.release(), "w");
   ZipWriter writer(zip_file);
   ASSERT_EQ(0, writer.StartEntry("compatibility.zip", kCompressStored));
   std::string compatibility_zip_content;
@@ -198,10 +199,10 @@ TEST(InstallTest, verify_package_compatibility_with_libvintf_system_manifest_xml
   CloseArchive(zip);
 }
 
-TEST(InstallTest, update_binary_command_smoke) {
 #ifdef AB_OTA_UPDATER
+static void VerifyAbUpdateBinaryCommand(const std::string& serialno, bool success = true) {
   TemporaryFile temp_file;
-  FILE* zip_file = fdopen(temp_file.fd, "w");
+  FILE* zip_file = fdopen(temp_file.release(), "w");
   ZipWriter writer(zip_file);
   ASSERT_EQ(0, writer.StartEntry("payload.bin", kCompressStored));
   ASSERT_EQ(0, writer.FinishEntry());
@@ -215,11 +216,13 @@ TEST(InstallTest, update_binary_command_smoke) {
   ASSERT_NE("", device);
   std::string timestamp = android::base::GetProperty("ro.build.date.utc", "");
   ASSERT_NE("", timestamp);
-  std::string metadata = android::base::Join(
-      std::vector<std::string>{
-          "ota-type=AB", "pre-device=" + device, "post-timestamp=" + timestamp,
-      },
-      "\n");
+
+  std::vector<std::string> meta{ "ota-type=AB", "pre-device=" + device,
+                                 "post-timestamp=" + timestamp };
+  if (!serialno.empty()) {
+    meta.push_back("serialno=" + serialno);
+  }
+  std::string metadata = android::base::Join(meta, "\n");
   ASSERT_EQ(0, writer.WriteBytes(metadata.data(), metadata.size()));
   ASSERT_EQ(0, writer.FinishEntry());
   ASSERT_EQ(0, writer.Finish());
@@ -234,17 +237,28 @@ TEST(InstallTest, update_binary_command_smoke) {
   std::string package = "/path/to/update.zip";
   std::string binary_path = "/sbin/update_engine_sideload";
   std::vector<std::string> cmd;
-  ASSERT_EQ(0, update_binary_command(package, zip, binary_path, 0, status_fd, &cmd));
-  ASSERT_EQ(5U, cmd.size());
-  ASSERT_EQ(binary_path, cmd[0]);
-  ASSERT_EQ("--payload=file://" + package, cmd[1]);
-  ASSERT_EQ("--offset=" + std::to_string(payload_entry.offset), cmd[2]);
-  ASSERT_EQ("--headers=" + properties, cmd[3]);
-  ASSERT_EQ("--status_fd=" + std::to_string(status_fd), cmd[4]);
+  if (success) {
+    ASSERT_EQ(0, update_binary_command(package, zip, binary_path, 0, status_fd, &cmd));
+    ASSERT_EQ(5U, cmd.size());
+    ASSERT_EQ(binary_path, cmd[0]);
+    ASSERT_EQ("--payload=file://" + package, cmd[1]);
+    ASSERT_EQ("--offset=" + std::to_string(payload_entry.offset), cmd[2]);
+    ASSERT_EQ("--headers=" + properties, cmd[3]);
+    ASSERT_EQ("--status_fd=" + std::to_string(status_fd), cmd[4]);
+  } else {
+    ASSERT_EQ(INSTALL_ERROR, update_binary_command(package, zip, binary_path, 0, status_fd, &cmd));
+  }
   CloseArchive(zip);
+}
+#endif  // AB_OTA_UPDATER
+
+TEST(InstallTest, update_binary_command_smoke) {
+#ifdef AB_OTA_UPDATER
+  // Empty serialno will pass the verification.
+  VerifyAbUpdateBinaryCommand({});
 #else
   TemporaryFile temp_file;
-  FILE* zip_file = fdopen(temp_file.fd, "w");
+  FILE* zip_file = fdopen(temp_file.release(), "w");
   ZipWriter writer(zip_file);
   static constexpr const char* UPDATE_BINARY_NAME = "META-INF/com/google/android/update-binary";
   ASSERT_EQ(0, writer.StartEntry(UPDATE_BINARY_NAME, kCompressStored));
@@ -289,7 +303,7 @@ TEST(InstallTest, update_binary_command_smoke) {
 TEST(InstallTest, update_binary_command_invalid) {
 #ifdef AB_OTA_UPDATER
   TemporaryFile temp_file;
-  FILE* zip_file = fdopen(temp_file.fd, "w");
+  FILE* zip_file = fdopen(temp_file.release(), "w");
   ZipWriter writer(zip_file);
   // Missing payload_properties.txt.
   ASSERT_EQ(0, writer.StartEntry("payload.bin", kCompressStored));
@@ -320,7 +334,7 @@ TEST(InstallTest, update_binary_command_invalid) {
   CloseArchive(zip);
 #else
   TemporaryFile temp_file;
-  FILE* zip_file = fdopen(temp_file.fd, "w");
+  FILE* zip_file = fdopen(temp_file.release(), "w");
   ZipWriter writer(zip_file);
   // The archive must have something to be opened correctly.
   ASSERT_EQ(0, writer.StartEntry("dummy_entry", 0));
@@ -340,3 +354,34 @@ TEST(InstallTest, update_binary_command_invalid) {
   CloseArchive(zip);
 #endif  // AB_OTA_UPDATER
 }
+
+#ifdef AB_OTA_UPDATER
+TEST(InstallTest, update_binary_command_multiple_serialno) {
+  std::string serialno = android::base::GetProperty("ro.serialno", "");
+  ASSERT_NE("", serialno);
+
+  // Single matching serialno will pass the verification.
+  VerifyAbUpdateBinaryCommand(serialno);
+
+  static constexpr char alphabet[] =
+      "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
+  auto generator = []() { return alphabet[rand() % (sizeof(alphabet) - 1)]; };
+
+  // Generate 900 random serial numbers.
+  std::string random_serial;
+  for (size_t i = 0; i < 900; i++) {
+    generate_n(back_inserter(random_serial), serialno.size(), generator);
+    random_serial.append("|");
+  }
+  // Random serialnos should fail the verification.
+  VerifyAbUpdateBinaryCommand(random_serial, false);
+
+  std::string long_serial = random_serial + serialno + "|";
+  for (size_t i = 0; i < 99; i++) {
+    generate_n(back_inserter(long_serial), serialno.size(), generator);
+    long_serial.append("|");
+  }
+  // String with the matching serialno should pass the verification.
+  VerifyAbUpdateBinaryCommand(long_serial);
+}
+#endif  // AB_OTA_UPDATER
diff --git a/tests/component/sideload_test.cpp b/tests/component/sideload_test.cpp
index 40cfc6975..b7109fcc2 100644
--- a/tests/component/sideload_test.cpp
+++ b/tests/component/sideload_test.cpp
@@ -16,6 +16,12 @@
 
 #include <unistd.h>
 
+#include <string>
+#include <vector>
+
+#include <android-base/file.h>
+#include <android-base/strings.h>
+#include <android-base/test_utils.h>
 #include <gtest/gtest.h>
 
 #include "fuse_sideload.h"
@@ -26,11 +32,67 @@ TEST(SideloadTest, fuse_device) {
 
 TEST(SideloadTest, run_fuse_sideload_wrong_parameters) {
   provider_vtab vtab;
-  vtab.close = [](void*) {};
+  vtab.close = [](void) {};
 
-  ASSERT_EQ(-1, run_fuse_sideload(&vtab, nullptr, 4096, 4095));
-  ASSERT_EQ(-1, run_fuse_sideload(&vtab, nullptr, 4096, (1 << 22) + 1));
+  ASSERT_EQ(-1, run_fuse_sideload(vtab, 4096, 4095));
+  ASSERT_EQ(-1, run_fuse_sideload(vtab, 4096, (1 << 22) + 1));
 
   // Too many blocks.
-  ASSERT_EQ(-1, run_fuse_sideload(&vtab, nullptr, ((1 << 18) + 1) * 4096, 4096));
+  ASSERT_EQ(-1, run_fuse_sideload(vtab, ((1 << 18) + 1) * 4096, 4096));
+}
+
+TEST(SideloadTest, run_fuse_sideload) {
+  const std::vector<std::string> blocks = {
+    std::string(2048, 'a') + std::string(2048, 'b'),
+    std::string(2048, 'c') + std::string(2048, 'd'),
+    std::string(2048, 'e') + std::string(2048, 'f'),
+    std::string(2048, 'g') + std::string(2048, 'h'),
+  };
+  const std::string content = android::base::Join(blocks, "");
+  ASSERT_EQ(16384U, content.size());
+
+  provider_vtab vtab;
+  vtab.close = [](void) {};
+  vtab.read_block = [&blocks](uint32_t block, uint8_t* buffer, uint32_t fetch_size) {
+    if (block >= 4) return -1;
+    blocks[block].copy(reinterpret_cast<char*>(buffer), fetch_size);
+    return 0;
+  };
+
+  TemporaryDir mount_point;
+  pid_t pid = fork();
+  if (pid == 0) {
+    ASSERT_EQ(0, run_fuse_sideload(vtab, 16384, 4096, mount_point.path));
+    _exit(EXIT_SUCCESS);
+  }
+
+  std::string package = std::string(mount_point.path) + "/" + FUSE_SIDELOAD_HOST_FILENAME;
+  int status;
+  static constexpr int kSideloadInstallTimeout = 10;
+  for (int i = 0; i < kSideloadInstallTimeout; ++i) {
+    ASSERT_NE(-1, waitpid(pid, &status, WNOHANG));
+
+    struct stat sb;
+    if (stat(package.c_str(), &sb) == 0) {
+      break;
+    }
+
+    if (errno == ENOENT && i < kSideloadInstallTimeout - 1) {
+      sleep(1);
+      continue;
+    }
+    FAIL() << "Timed out waiting for the fuse-provided package.";
+  }
+
+  std::string content_via_fuse;
+  ASSERT_TRUE(android::base::ReadFileToString(package, &content_via_fuse));
+  ASSERT_EQ(content, content_via_fuse);
+
+  std::string exit_flag = std::string(mount_point.path) + "/" + FUSE_SIDELOAD_HOST_EXIT_FLAG;
+  struct stat sb;
+  ASSERT_EQ(0, stat(exit_flag.c_str(), &sb));
+
+  waitpid(pid, &status, 0);
+  ASSERT_EQ(0, WEXITSTATUS(status));
+  ASSERT_EQ(EXIT_SUCCESS, WEXITSTATUS(status));
 }
diff --git a/tests/component/uncrypt_test.cpp b/tests/component/uncrypt_test.cpp
index 3925236a5..55baca2e3 100644
--- a/tests/component/uncrypt_test.cpp
+++ b/tests/component/uncrypt_test.cpp
@@ -20,6 +20,7 @@
 #include <sys/un.h>
 #include <unistd.h>
 
+#include <algorithm>
 #include <string>
 
 #include <android-base/file.h>
@@ -38,43 +39,49 @@ static const std::string INIT_SVC_CLEAR_BCB = "init.svc.clear-bcb";
 static const std::string INIT_SVC_UNCRYPT = "init.svc.uncrypt";
 static constexpr int SOCKET_CONNECTION_MAX_RETRY = 30;
 
+static void StopService() {
+  ASSERT_TRUE(android::base::SetProperty("ctl.stop", "setup-bcb"));
+  ASSERT_TRUE(android::base::SetProperty("ctl.stop", "clear-bcb"));
+  ASSERT_TRUE(android::base::SetProperty("ctl.stop", "uncrypt"));
+
+  bool success = false;
+  for (int retry = 0; retry < SOCKET_CONNECTION_MAX_RETRY; retry++) {
+    std::string setup_bcb = android::base::GetProperty(INIT_SVC_SETUP_BCB, "");
+    std::string clear_bcb = android::base::GetProperty(INIT_SVC_CLEAR_BCB, "");
+    std::string uncrypt = android::base::GetProperty(INIT_SVC_UNCRYPT, "");
+    GTEST_LOG_(INFO) << "setup-bcb: [" << setup_bcb << "] clear-bcb: [" << clear_bcb
+                     << "] uncrypt: [" << uncrypt << "]";
+    if (setup_bcb != "running" && clear_bcb != "running" && uncrypt != "running") {
+      success = true;
+      break;
+    }
+    sleep(1);
+  }
+
+  ASSERT_TRUE(success) << "uncrypt service is not available.";
+}
+
 class UncryptTest : public ::testing::Test {
  protected:
   UncryptTest() : has_misc(true) {}
 
-  virtual void SetUp() override {
-    ASSERT_TRUE(android::base::SetProperty("ctl.stop", "setup-bcb"));
-    ASSERT_TRUE(android::base::SetProperty("ctl.stop", "clear-bcb"));
-    ASSERT_TRUE(android::base::SetProperty("ctl.stop", "uncrypt"));
-
-    bool success = false;
-    for (int retry = 0; retry < SOCKET_CONNECTION_MAX_RETRY; retry++) {
-      std::string setup_bcb = android::base::GetProperty(INIT_SVC_SETUP_BCB, "");
-      std::string clear_bcb = android::base::GetProperty(INIT_SVC_CLEAR_BCB, "");
-      std::string uncrypt = android::base::GetProperty(INIT_SVC_UNCRYPT, "");
-      LOG(INFO) << "setup-bcb: [" << setup_bcb << "] clear-bcb: [" << clear_bcb << "] uncrypt: ["
-                << uncrypt << "]";
-      if (setup_bcb != "running" && clear_bcb != "running" && uncrypt != "running") {
-        success = true;
-        break;
-      }
-      sleep(1);
-    }
-
-    ASSERT_TRUE(success) << "uncrypt service is not available.";
-
+  void SetUp() override {
     std::string err;
     has_misc = !get_bootloader_message_blk_device(&err).empty();
   }
 
-  void SetupOrClearBcb(bool isSetup, const std::string& message,
-                       const std::string& message_in_bcb) const {
-    if (!has_misc) {
-      GTEST_LOG_(INFO) << "Test skipped due to no /misc partition found on the device.";
-      return;
+  void TearDown() override {
+    // Clear the BCB.
+    if (has_misc) {
+      std::string err;
+      ASSERT_TRUE(clear_bootloader_message(&err)) << "Failed to clear BCB: " << err;
     }
+  }
 
-    // Trigger the setup-bcb service.
+  void SetupOrClearBcb(bool isSetup, const std::string& message,
+                       const std::string& message_in_bcb) const {
+    // Restart the setup-bcb service.
+    StopService();
     ASSERT_TRUE(android::base::SetProperty("ctl.start", isSetup ? "setup-bcb" : "clear-bcb"));
 
     // Test tends to be flaky if proceeding immediately ("Transport endpoint is not connected").
@@ -144,27 +151,49 @@ class UncryptTest : public ::testing::Test {
     }
   }
 
+  void VerifyBootloaderMessage(const std::string& expected) {
+    std::string err;
+    bootloader_message boot;
+    ASSERT_TRUE(read_bootloader_message(&boot, &err)) << "Failed to read BCB: " << err;
+
+    // Check that we have all the expected bytes.
+    ASSERT_EQ(expected, std::string(reinterpret_cast<const char*>(&boot), sizeof(boot)));
+  }
+
   bool has_misc;
 };
 
 TEST_F(UncryptTest, setup_bcb) {
+  if (!has_misc) {
+    GTEST_LOG_(INFO) << "Test skipped due to no /misc partition found on the device.";
+    return;
+  }
+
+  std::string random_data;
+  random_data.reserve(sizeof(bootloader_message));
+  generate_n(back_inserter(random_data), sizeof(bootloader_message), []() { return rand() % 128; });
+
+  bootloader_message boot;
+  memcpy(&boot, random_data.c_str(), random_data.size());
+
+  std::string err;
+  ASSERT_TRUE(write_bootloader_message(boot, &err)) << "Failed to write BCB: " << err;
+  VerifyBootloaderMessage(random_data);
+
+  ASSERT_TRUE(clear_bootloader_message(&err)) << "Failed to clear BCB: " << err;
+  VerifyBootloaderMessage(std::string(sizeof(bootloader_message), '\0'));
+
   std::string message = "--update_message=abc value";
   std::string message_in_bcb = "recovery\n--update_message=abc value\n";
   SetupOrClearBcb(true, message, message_in_bcb);
-}
 
-TEST_F(UncryptTest, clear_bcb) {
   SetupOrClearBcb(false, "", "");
-}
 
-TEST_F(UncryptTest, setup_bcb_wipe_ab) {
   TemporaryFile wipe_package;
   ASSERT_TRUE(android::base::WriteStringToFile(std::string(345, 'a'), wipe_package.path));
 
   // It's expected to store a wipe package in /misc, with the package size passed to recovery.
-  std::string message =
-      "--wipe_ab\n--wipe_package="s + wipe_package.path + "\n--reason=wipePackage"s;
-  std::string message_in_bcb =
-      "recovery\n--wipe_ab\n--wipe_package_size=345\n--reason=wipePackage\n";
+  message = "--wipe_ab\n--wipe_package="s + wipe_package.path + "\n--reason=wipePackage"s;
+  message_in_bcb = "recovery\n--wipe_ab\n--wipe_package_size=345\n--reason=wipePackage\n";
   SetupOrClearBcb(true, message, message_in_bcb);
 }
diff --git a/tests/component/update_verifier_test.cpp b/tests/component/update_verifier_test.cpp
index b04e1185e..1544bb2a4 100644
--- a/tests/component/update_verifier_test.cpp
+++ b/tests/component/update_verifier_test.cpp
@@ -46,7 +46,6 @@ TEST_F(UpdateVerifierTest, verify_image_smoke) {
     return;
   }
 
-  // The care map file can have only two or four lines.
   TemporaryFile temp_file;
   std::string content = "system\n2,0,1";
   ASSERT_TRUE(android::base::WriteStringToFile(content, temp_file.path));
@@ -58,7 +57,7 @@ TEST_F(UpdateVerifierTest, verify_image_smoke) {
 }
 
 TEST_F(UpdateVerifierTest, verify_image_wrong_lines) {
-  // The care map file can have only two or four lines.
+  // The care map file can have only 2 / 4 / 6 lines.
   TemporaryFile temp_file;
   ASSERT_FALSE(verify_image(temp_file.path));
 
diff --git a/tests/component/updater_test.cpp b/tests/component/updater_test.cpp
index 374a99c89..50e0a6345 100644
--- a/tests/component/updater_test.cpp
+++ b/tests/component/updater_test.cpp
@@ -23,6 +23,7 @@
 #include <algorithm>
 #include <memory>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 #include <android-base/file.h>
@@ -32,16 +33,17 @@
 #include <android-base/test_utils.h>
 #include <bootloader_message/bootloader_message.h>
 #include <brotli/encode.h>
-#include <bsdiff.h>
+#include <bsdiff/bsdiff.h>
 #include <gtest/gtest.h>
 #include <ziparchive/zip_archive.h>
 #include <ziparchive/zip_writer.h>
 
 #include "common/test_constants.h"
 #include "edify/expr.h"
-#include "error_code.h"
 #include "otautil/SysUtil.h"
-#include "print_sha1.h"
+#include "otautil/cache_location.h"
+#include "otautil/error_code.h"
+#include "otautil/print_sha1.h"
 #include "updater/blockimg.h"
 #include "updater/install.h"
 #include "updater/updater.h"
@@ -74,6 +76,23 @@ static void expect(const char* expected, const char* expr_str, CauseCode cause_c
   ASSERT_EQ(cause_code, state.cause_code);
 }
 
+static void BuildUpdatePackage(const std::unordered_map<std::string, std::string>& entries,
+                               int fd) {
+  FILE* zip_file_ptr = fdopen(fd, "wb");
+  ZipWriter zip_writer(zip_file_ptr);
+
+  for (const auto& entry : entries) {
+    ASSERT_EQ(0, zip_writer.StartEntry(entry.first.c_str(), 0));
+    if (!entry.second.empty()) {
+      ASSERT_EQ(0, zip_writer.WriteBytes(entry.second.data(), entry.second.size()));
+    }
+    ASSERT_EQ(0, zip_writer.FinishEntry());
+  }
+
+  ASSERT_EQ(0, zip_writer.Finish());
+  ASSERT_EQ(0, fclose(zip_file_ptr));
+}
+
 static std::string get_sha1(const std::string& content) {
   uint8_t digest[SHA_DIGEST_LENGTH];
   SHA1(reinterpret_cast<const uint8_t*>(content.c_str()), content.size(), digest);
@@ -86,7 +105,16 @@ class UpdaterTest : public ::testing::Test {
     RegisterBuiltins();
     RegisterInstallFunctions();
     RegisterBlockImageFunctions();
+
+    // Mock the location of last_command_file.
+    CacheLocation::location().set_cache_temp_source(temp_saved_source_.path);
+    CacheLocation::location().set_last_command_file(temp_last_command_.path);
+    CacheLocation::location().set_stash_directory_base(temp_stash_base_.path);
   }
+
+  TemporaryFile temp_saved_source_;
+  TemporaryFile temp_last_command_;
+  TemporaryDir temp_stash_base_;
 };
 
 TEST_F(UpdaterTest, getprop) {
@@ -592,7 +620,7 @@ TEST_F(UpdaterTest, set_progress) {
 
   TemporaryFile tf;
   UpdaterInfo updater_info;
-  updater_info.cmd_pipe = fdopen(tf.fd, "w");
+  updater_info.cmd_pipe = fdopen(tf.release(), "w");
   expect(".52", "set_progress(\".52\")", kNoCause, &updater_info);
   fflush(updater_info.cmd_pipe);
 
@@ -601,6 +629,7 @@ TEST_F(UpdaterTest, set_progress) {
   ASSERT_EQ(android::base::StringPrintf("set_progress %f\n", .52), cmd);
   // recovery-updater protocol expects 2 tokens ("set_progress <frac>").
   ASSERT_EQ(2U, android::base::Split(cmd, " ").size());
+  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
 }
 
 TEST_F(UpdaterTest, show_progress) {
@@ -616,7 +645,7 @@ TEST_F(UpdaterTest, show_progress) {
 
   TemporaryFile tf;
   UpdaterInfo updater_info;
-  updater_info.cmd_pipe = fdopen(tf.fd, "w");
+  updater_info.cmd_pipe = fdopen(tf.release(), "w");
   expect(".52", "show_progress(\".52\", \"10\")", kNoCause, &updater_info);
   fflush(updater_info.cmd_pipe);
 
@@ -625,32 +654,22 @@ TEST_F(UpdaterTest, show_progress) {
   ASSERT_EQ(android::base::StringPrintf("progress %f %d\n", .52, 10), cmd);
   // recovery-updater protocol expects 3 tokens ("progress <frac> <secs>").
   ASSERT_EQ(3U, android::base::Split(cmd, " ").size());
+  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
 }
 
-TEST_F(UpdaterTest, block_image_update) {
-  // Create a zip file with new_data and patch_data.
-  TemporaryFile zip_file;
-  FILE* zip_file_ptr = fdopen(zip_file.fd, "wb");
-  ZipWriter zip_writer(zip_file_ptr);
-
-  // Add a dummy new data.
-  ASSERT_EQ(0, zip_writer.StartEntry("new_data", 0));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-
-  // Generate and add the patch data.
+TEST_F(UpdaterTest, block_image_update_patch_data) {
   std::string src_content = std::string(4096, 'a') + std::string(4096, 'c');
   std::string tgt_content = std::string(4096, 'b') + std::string(4096, 'd');
+
+  // Generate the patch data.
   TemporaryFile patch_file;
   ASSERT_EQ(0, bsdiff::bsdiff(reinterpret_cast<const uint8_t*>(src_content.data()),
       src_content.size(), reinterpret_cast<const uint8_t*>(tgt_content.data()),
       tgt_content.size(), patch_file.path, nullptr));
   std::string patch_content;
   ASSERT_TRUE(android::base::ReadFileToString(patch_file.path, &patch_content));
-  ASSERT_EQ(0, zip_writer.StartEntry("patch_data", 0));
-  ASSERT_EQ(0, zip_writer.WriteBytes(patch_content.data(), patch_content.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
 
-  // Add two transfer lists. The first one contains a bsdiff; and we expect the update to succeed.
+  // Create the transfer list that contains a bsdiff.
   std::string src_hash = get_sha1(src_content);
   std::string tgt_hash = get_sha1(tgt_content);
   std::vector<std::string> transfer_list = {
@@ -663,27 +682,16 @@ TEST_F(UpdaterTest, block_image_update) {
                                 src_hash.c_str(), tgt_hash.c_str(), src_hash.c_str()),
     "free " + src_hash,
   };
-  ASSERT_EQ(0, zip_writer.StartEntry("transfer_list", 0));
-  std::string commands = android::base::Join(transfer_list, '\n');
-  ASSERT_EQ(0, zip_writer.WriteBytes(commands.data(), commands.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
 
-  // Stash and free some blocks, then fail the 2nd update intentionally.
-  std::vector<std::string> fail_transfer_list = {
-    "4",
-    "2",
-    "0",
-    "2",
-    "stash " + tgt_hash + " 2,0,2",
-    "free " + tgt_hash,
-    "fail",
+  std::unordered_map<std::string, std::string> entries = {
+    { "new_data", "" },
+    { "patch_data", patch_content },
+    { "transfer_list", android::base::Join(transfer_list, '\n') },
   };
-  ASSERT_EQ(0, zip_writer.StartEntry("fail_transfer_list", 0));
-  std::string fail_commands = android::base::Join(fail_transfer_list, '\n');
-  ASSERT_EQ(0, zip_writer.WriteBytes(fail_commands.data(), fail_commands.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  ASSERT_EQ(0, zip_writer.Finish());
-  ASSERT_EQ(0, fclose(zip_file_ptr));
+
+  // Build the update package.
+  TemporaryFile zip_file;
+  BuildUpdatePackage(entries, zip_file.release());
 
   MemMapping map;
   ASSERT_TRUE(map.MapFile(zip_file.path));
@@ -694,11 +702,11 @@ TEST_F(UpdaterTest, block_image_update) {
   UpdaterInfo updater_info;
   updater_info.package_zip = handle;
   TemporaryFile temp_pipe;
-  updater_info.cmd_pipe = fopen(temp_pipe.path, "wbe");
+  updater_info.cmd_pipe = fdopen(temp_pipe.release(), "wbe");
   updater_info.package_zip_addr = map.addr;
   updater_info.package_zip_len = map.length;
 
-  // Execute the commands in the 1st transfer list.
+  // Execute the commands in the transfer list.
   TemporaryFile update_file;
   ASSERT_TRUE(android::base::WriteStringToFile(src_content, update_file.path));
   std::string script = "block_image_update(\"" + std::string(update_file.path) +
@@ -709,44 +717,98 @@ TEST_F(UpdaterTest, block_image_update) {
   ASSERT_TRUE(android::base::ReadFileToString(update_file.path, &updated_content));
   ASSERT_EQ(tgt_hash, get_sha1(updated_content));
 
-  // Expect the 2nd update to fail, but expect the stashed blocks to be freed.
-  script = "block_image_update(\"" + std::string(update_file.path) +
-      R"(", package_extract_file("fail_transfer_list"), "new_data", "patch_data"))";
+  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
+  CloseArchive(handle);
+}
+
+TEST_F(UpdaterTest, block_image_update_fail) {
+  std::string src_content(4096 * 2, 'e');
+  std::string src_hash = get_sha1(src_content);
+  // Stash and free some blocks, then fail the update intentionally.
+  std::vector<std::string> transfer_list = {
+    "4", "2", "0", "2", "stash " + src_hash + " 2,0,2", "free " + src_hash, "fail",
+  };
+
+  // Add a new data of 10 bytes to test the deadlock.
+  std::unordered_map<std::string, std::string> entries = {
+    { "new_data", std::string(10, 0) },
+    { "patch_data", "" },
+    { "transfer_list", android::base::Join(transfer_list, '\n') },
+  };
+
+  // Build the update package.
+  TemporaryFile zip_file;
+  BuildUpdatePackage(entries, zip_file.release());
+
+  MemMapping map;
+  ASSERT_TRUE(map.MapFile(zip_file.path));
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));
+
+  // Set up the handler, command_pipe, patch offset & length.
+  UpdaterInfo updater_info;
+  updater_info.package_zip = handle;
+  TemporaryFile temp_pipe;
+  updater_info.cmd_pipe = fdopen(temp_pipe.release(), "wbe");
+  updater_info.package_zip_addr = map.addr;
+  updater_info.package_zip_len = map.length;
+
+  TemporaryFile update_file;
+  ASSERT_TRUE(android::base::WriteStringToFile(src_content, update_file.path));
+  // Expect the stashed blocks to be freed.
+  std::string script = "block_image_update(\"" + std::string(update_file.path) +
+                       R"(", package_extract_file("transfer_list"), "new_data", "patch_data"))";
   expect("", script.c_str(), kNoCause, &updater_info);
   // Updater generates the stash name based on the input file name.
   std::string name_digest = get_sha1(update_file.path);
-  std::string stash_base = "/cache/recovery/" + name_digest;
+  std::string stash_base = std::string(temp_stash_base_.path) + "/" + name_digest;
   ASSERT_EQ(0, access(stash_base.c_str(), F_OK));
-  ASSERT_EQ(-1, access((stash_base + tgt_hash).c_str(), F_OK));
+  ASSERT_EQ(-1, access((stash_base + src_hash).c_str(), F_OK));
   ASSERT_EQ(0, rmdir(stash_base.c_str()));
 
   ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
   CloseArchive(handle);
 }
 
-TEST_F(UpdaterTest, new_data_short_write) {
-  // Create a zip file with new_data.
+TEST_F(UpdaterTest, new_data_over_write) {
+  std::vector<std::string> transfer_list = {
+    "4", "1", "0", "0", "new 2,0,1",
+  };
+
+  // Write 4096 + 100 bytes of new data.
+  std::unordered_map<std::string, std::string> entries = {
+    { "new_data", std::string(4196, 0) },
+    { "patch_data", "" },
+    { "transfer_list", android::base::Join(transfer_list, '\n') },
+  };
+
+  // Build the update package.
   TemporaryFile zip_file;
-  FILE* zip_file_ptr = fdopen(zip_file.fd, "wb");
-  ZipWriter zip_writer(zip_file_ptr);
+  BuildUpdatePackage(entries, zip_file.release());
 
-  // Add the empty new data.
-  ASSERT_EQ(0, zip_writer.StartEntry("empty_new_data", 0));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  // Add the short written new data.
-  ASSERT_EQ(0, zip_writer.StartEntry("short_new_data", 0));
-  std::string new_data_short = std::string(10, 'a');
-  ASSERT_EQ(0, zip_writer.WriteBytes(new_data_short.data(), new_data_short.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  // Add the data of exactly one block.
-  ASSERT_EQ(0, zip_writer.StartEntry("exact_new_data", 0));
-  std::string new_data_exact = std::string(4096, 'a');
-  ASSERT_EQ(0, zip_writer.WriteBytes(new_data_exact.data(), new_data_exact.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  // Add a dummy patch data.
-  ASSERT_EQ(0, zip_writer.StartEntry("patch_data", 0));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
+  MemMapping map;
+  ASSERT_TRUE(map.MapFile(zip_file.path));
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));
 
+  // Set up the handler, command_pipe, patch offset & length.
+  UpdaterInfo updater_info;
+  updater_info.package_zip = handle;
+  TemporaryFile temp_pipe;
+  updater_info.cmd_pipe = fdopen(temp_pipe.release(), "wbe");
+  updater_info.package_zip_addr = map.addr;
+  updater_info.package_zip_len = map.length;
+
+  TemporaryFile update_file;
+  std::string script = "block_image_update(\"" + std::string(update_file.path) +
+                       R"(", package_extract_file("transfer_list"), "new_data", "patch_data"))";
+  expect("t", script.c_str(), kNoCause, &updater_info);
+
+  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
+  CloseArchive(handle);
+}
+
+TEST_F(UpdaterTest, new_data_short_write) {
   std::vector<std::string> transfer_list = {
     "4",
     "1",
@@ -754,12 +816,17 @@ TEST_F(UpdaterTest, new_data_short_write) {
     "0",
     "new 2,0,1",
   };
-  ASSERT_EQ(0, zip_writer.StartEntry("transfer_list", 0));
-  std::string commands = android::base::Join(transfer_list, '\n');
-  ASSERT_EQ(0, zip_writer.WriteBytes(commands.data(), commands.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  ASSERT_EQ(0, zip_writer.Finish());
-  ASSERT_EQ(0, fclose(zip_file_ptr));
+
+  std::unordered_map<std::string, std::string> entries = {
+    { "empty_new_data", "" },
+    { "short_new_data", std::string(10, 'a') },
+    { "exact_new_data", std::string(4096, 'a') },
+    { "patch_data", "" },
+    { "transfer_list", android::base::Join(transfer_list, '\n') },
+  };
+
+  TemporaryFile zip_file;
+  BuildUpdatePackage(entries, zip_file.release());
 
   MemMapping map;
   ASSERT_TRUE(map.MapFile(zip_file.path));
@@ -770,7 +837,7 @@ TEST_F(UpdaterTest, new_data_short_write) {
   UpdaterInfo updater_info;
   updater_info.package_zip = handle;
   TemporaryFile temp_pipe;
-  updater_info.cmd_pipe = fopen(temp_pipe.path, "wbe");
+  updater_info.cmd_pipe = fdopen(temp_pipe.release(), "wbe");
   updater_info.package_zip_addr = map.addr;
   updater_info.package_zip_len = map.length;
 
@@ -788,18 +855,12 @@ TEST_F(UpdaterTest, new_data_short_write) {
   std::string script_exact_data = "block_image_update(\"" + std::string(update_file.path) +
       R"(", package_extract_file("transfer_list"), "exact_new_data", "patch_data"))";
   expect("t", script_exact_data.c_str(), kNoCause, &updater_info);
+
+  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
   CloseArchive(handle);
 }
 
 TEST_F(UpdaterTest, brotli_new_data) {
-  // Create a zip file with new_data.
-  TemporaryFile zip_file;
-  FILE* zip_file_ptr = fdopen(zip_file.fd, "wb");
-  ZipWriter zip_writer(zip_file_ptr);
-
-  // Add a brotli compressed new data entry.
-  ASSERT_EQ(0, zip_writer.StartEntry("new.dat.br", 0));
-
   auto generator = []() { return rand() % 128; };
   // Generate 100 blocks of random data.
   std::string brotli_new_data;
@@ -807,16 +868,12 @@ TEST_F(UpdaterTest, brotli_new_data) {
   generate_n(back_inserter(brotli_new_data), 4096 * 100, generator);
 
   size_t encoded_size = BrotliEncoderMaxCompressedSize(brotli_new_data.size());
-  std::vector<uint8_t> encoded_data(encoded_size);
+  std::string encoded_data(encoded_size, 0);
   ASSERT_TRUE(BrotliEncoderCompress(
       BROTLI_DEFAULT_QUALITY, BROTLI_DEFAULT_WINDOW, BROTLI_DEFAULT_MODE, brotli_new_data.size(),
-      reinterpret_cast<const uint8_t*>(brotli_new_data.data()), &encoded_size, encoded_data.data()));
-
-  ASSERT_EQ(0, zip_writer.WriteBytes(encoded_data.data(), encoded_size));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  // Add a dummy patch data.
-  ASSERT_EQ(0, zip_writer.StartEntry("patch_data", 0));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
+      reinterpret_cast<const uint8_t*>(brotli_new_data.data()), &encoded_size,
+      reinterpret_cast<uint8_t*>(const_cast<char*>(encoded_data.data()))));
+  encoded_data.resize(encoded_size);
 
   // Write a few small chunks of new data, then a large chunk, and finally a few small chunks.
   // This helps us to catch potential short writes.
@@ -832,12 +889,15 @@ TEST_F(UpdaterTest, brotli_new_data) {
     "new 2,98,99",
     "new 2,99,100",
   };
-  ASSERT_EQ(0, zip_writer.StartEntry("transfer_list", 0));
-  std::string commands = android::base::Join(transfer_list, '\n');
-  ASSERT_EQ(0, zip_writer.WriteBytes(commands.data(), commands.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  ASSERT_EQ(0, zip_writer.Finish());
-  ASSERT_EQ(0, fclose(zip_file_ptr));
+
+  std::unordered_map<std::string, std::string> entries = {
+    { "new.dat.br", std::move(encoded_data) },
+    { "patch_data", "" },
+    { "transfer_list", android::base::Join(transfer_list, '\n') },
+  };
+
+  TemporaryFile zip_file;
+  BuildUpdatePackage(entries, zip_file.release());
 
   MemMapping map;
   ASSERT_TRUE(map.MapFile(zip_file.path));
@@ -848,7 +908,7 @@ TEST_F(UpdaterTest, brotli_new_data) {
   UpdaterInfo updater_info;
   updater_info.package_zip = handle;
   TemporaryFile temp_pipe;
-  updater_info.cmd_pipe = fopen(temp_pipe.path, "wb");
+  updater_info.cmd_pipe = fdopen(temp_pipe.release(), "wb");
   updater_info.package_zip_addr = map.addr;
   updater_info.package_zip_len = map.length;
 
@@ -862,5 +922,221 @@ TEST_F(UpdaterTest, brotli_new_data) {
   std::string updated_content;
   ASSERT_TRUE(android::base::ReadFileToString(update_file.path, &updated_content));
   ASSERT_EQ(brotli_new_data, updated_content);
+
+  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
+  CloseArchive(handle);
+}
+
+TEST_F(UpdaterTest, last_command_update) {
+  std::string last_command_file = CacheLocation::location().last_command_file();
+
+  std::string block1 = std::string(4096, '1');
+  std::string block2 = std::string(4096, '2');
+  std::string block3 = std::string(4096, '3');
+  std::string block1_hash = get_sha1(block1);
+  std::string block2_hash = get_sha1(block2);
+  std::string block3_hash = get_sha1(block3);
+
+  // Compose the transfer list to fail the first update.
+  std::vector<std::string> transfer_list_fail = {
+    "4",
+    "2",
+    "0",
+    "2",
+    "stash " + block1_hash + " 2,0,1",
+    "move " + block1_hash + " 2,1,2 1 2,0,1",
+    "stash " + block3_hash + " 2,2,3",
+    "fail",
+  };
+
+  // Mimic a resumed update with the same transfer commands.
+  std::vector<std::string> transfer_list_continue = {
+    "4",
+    "2",
+    "0",
+    "2",
+    "stash " + block1_hash + " 2,0,1",
+    "move " + block1_hash + " 2,1,2 1 2,0,1",
+    "stash " + block3_hash + " 2,2,3",
+    "move " + block1_hash + " 2,2,3 1 2,0,1",
+  };
+
+  std::unordered_map<std::string, std::string> entries = {
+    { "new_data", "" },
+    { "patch_data", "" },
+    { "transfer_list_fail", android::base::Join(transfer_list_fail, '\n') },
+    { "transfer_list_continue", android::base::Join(transfer_list_continue, '\n') },
+  };
+
+  // Build the update package.
+  TemporaryFile zip_file;
+  BuildUpdatePackage(entries, zip_file.release());
+
+  MemMapping map;
+  ASSERT_TRUE(map.MapFile(zip_file.path));
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));
+
+  // Set up the handler, command_pipe, patch offset & length.
+  UpdaterInfo updater_info;
+  updater_info.package_zip = handle;
+  TemporaryFile temp_pipe;
+  updater_info.cmd_pipe = fdopen(temp_pipe.release(), "wbe");
+  updater_info.package_zip_addr = map.addr;
+  updater_info.package_zip_len = map.length;
+
+  std::string src_content = block1 + block2 + block3;
+  TemporaryFile update_file;
+  ASSERT_TRUE(android::base::WriteStringToFile(src_content, update_file.path));
+  std::string script =
+      "block_image_update(\"" + std::string(update_file.path) +
+      R"(", package_extract_file("transfer_list_fail"), "new_data", "patch_data"))";
+  expect("", script.c_str(), kNoCause, &updater_info);
+
+  // Expect last_command to contain the last stash command.
+  std::string last_command_content;
+  ASSERT_TRUE(android::base::ReadFileToString(last_command_file.c_str(), &last_command_content));
+  EXPECT_EQ("2\nstash " + block3_hash + " 2,2,3", last_command_content);
+  std::string updated_contents;
+  ASSERT_TRUE(android::base::ReadFileToString(update_file.path, &updated_contents));
+  ASSERT_EQ(block1 + block1 + block3, updated_contents);
+
+  // Resume the update, expect the first 'move' to be skipped but the second 'move' to be executed.
+  ASSERT_TRUE(android::base::WriteStringToFile(src_content, update_file.path));
+  std::string script_second_update =
+      "block_image_update(\"" + std::string(update_file.path) +
+      R"(", package_extract_file("transfer_list_continue"), "new_data", "patch_data"))";
+  expect("t", script_second_update.c_str(), kNoCause, &updater_info);
+  ASSERT_TRUE(android::base::ReadFileToString(update_file.path, &updated_contents));
+  ASSERT_EQ(block1 + block2 + block1, updated_contents);
+
+  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
+  CloseArchive(handle);
+}
+
+TEST_F(UpdaterTest, last_command_update_unresumable) {
+  std::string last_command_file = CacheLocation::location().last_command_file();
+
+  std::string block1 = std::string(4096, '1');
+  std::string block2 = std::string(4096, '2');
+  std::string block1_hash = get_sha1(block1);
+  std::string block2_hash = get_sha1(block2);
+
+  // Construct an unresumable update with source blocks mismatch.
+  std::vector<std::string> transfer_list_unresumable = {
+    "4", "2", "0", "2", "stash " + block1_hash + " 2,0,1", "move " + block2_hash + " 2,1,2 1 2,0,1",
+  };
+
+  std::unordered_map<std::string, std::string> entries = {
+    { "new_data", "" },
+    { "patch_data", "" },
+    { "transfer_list_unresumable", android::base::Join(transfer_list_unresumable, '\n') },
+  };
+
+  // Build the update package.
+  TemporaryFile zip_file;
+  BuildUpdatePackage(entries, zip_file.release());
+
+  MemMapping map;
+  ASSERT_TRUE(map.MapFile(zip_file.path));
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));
+
+  // Set up the handler, command_pipe, patch offset & length.
+  UpdaterInfo updater_info;
+  updater_info.package_zip = handle;
+  TemporaryFile temp_pipe;
+  updater_info.cmd_pipe = fdopen(temp_pipe.release(), "wbe");
+  updater_info.package_zip_addr = map.addr;
+  updater_info.package_zip_len = map.length;
+
+  // Set up the last_command_file
+  ASSERT_TRUE(
+      android::base::WriteStringToFile("0\nstash " + block1_hash + " 2,0,1", last_command_file));
+
+  // The last_command_file will be deleted if the update encounters an unresumable failure
+  // later.
+  std::string src_content = block1 + block1;
+  TemporaryFile update_file;
+  ASSERT_TRUE(android::base::WriteStringToFile(src_content, update_file.path));
+  std::string script =
+      "block_image_update(\"" + std::string(update_file.path) +
+      R"(", package_extract_file("transfer_list_unresumable"), "new_data", "patch_data"))";
+  expect("", script.c_str(), kNoCause, &updater_info);
+  ASSERT_EQ(-1, access(last_command_file.c_str(), R_OK));
+
+  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
+  CloseArchive(handle);
+}
+
+TEST_F(UpdaterTest, last_command_verify) {
+  std::string last_command_file = CacheLocation::location().last_command_file();
+
+  std::string block1 = std::string(4096, '1');
+  std::string block2 = std::string(4096, '2');
+  std::string block3 = std::string(4096, '3');
+  std::string block1_hash = get_sha1(block1);
+  std::string block2_hash = get_sha1(block2);
+  std::string block3_hash = get_sha1(block3);
+
+  std::vector<std::string> transfer_list_verify = {
+    "4",
+    "2",
+    "0",
+    "2",
+    "stash " + block1_hash + " 2,0,1",
+    "move " + block1_hash + " 2,0,1 1 2,0,1",
+    "move " + block1_hash + " 2,1,2 1 2,0,1",
+    "stash " + block3_hash + " 2,2,3",
+  };
+
+  std::unordered_map<std::string, std::string> entries = {
+    { "new_data", "" },
+    { "patch_data", "" },
+    { "transfer_list_verify", android::base::Join(transfer_list_verify, '\n') },
+  };
+
+  // Build the update package.
+  TemporaryFile zip_file;
+  BuildUpdatePackage(entries, zip_file.release());
+
+  MemMapping map;
+  ASSERT_TRUE(map.MapFile(zip_file.path));
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));
+
+  // Set up the handler, command_pipe, patch offset & length.
+  UpdaterInfo updater_info;
+  updater_info.package_zip = handle;
+  TemporaryFile temp_pipe;
+  updater_info.cmd_pipe = fdopen(temp_pipe.release(), "wbe");
+  updater_info.package_zip_addr = map.addr;
+  updater_info.package_zip_len = map.length;
+
+  std::string src_content = block1 + block1 + block3;
+  TemporaryFile update_file;
+  ASSERT_TRUE(android::base::WriteStringToFile(src_content, update_file.path));
+
+  ASSERT_TRUE(
+      android::base::WriteStringToFile("2\nstash " + block3_hash + " 2,2,3", last_command_file));
+
+  // Expect the verification to succeed and the last_command_file is intact.
+  std::string script_verify =
+      "block_image_verify(\"" + std::string(update_file.path) +
+      R"(", package_extract_file("transfer_list_verify"), "new_data","patch_data"))";
+  expect("t", script_verify.c_str(), kNoCause, &updater_info);
+
+  std::string last_command_content;
+  ASSERT_TRUE(android::base::ReadFileToString(last_command_file.c_str(), &last_command_content));
+  EXPECT_EQ("2\nstash " + block3_hash + " 2,2,3", last_command_content);
+
+  // Expect the verification to succeed but last_command_file to be deleted; because the target
+  // blocks don't have the expected contents for the second move command.
+  src_content = block1 + block2 + block3;
+  ASSERT_TRUE(android::base::WriteStringToFile(src_content, update_file.path));
+  expect("t", script_verify.c_str(), kNoCause, &updater_info);
+  ASSERT_EQ(-1, access(last_command_file.c_str(), R_OK));
+
+  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
   CloseArchive(handle);
 }
diff --git a/tests/manual/recovery_test.cpp b/tests/manual/recovery_test.cpp
index 92c6ef2d4..64e3b59e6 100644
--- a/tests/manual/recovery_test.cpp
+++ b/tests/manual/recovery_test.cpp
@@ -209,7 +209,7 @@ TEST_P(ResourceTest, ValidateLocale) {
     ASSERT_GT(height, y + 1 + h) << "Locale: " << kLocale << " is not found in the file.";
     char* loc = reinterpret_cast<char*>(&row[5]);
     if (matches_locale(loc, kLocale.c_str())) {
-      EXPECT_TRUE(android::base::StartsWith(loc, kLocale.c_str()));
+      EXPECT_TRUE(android::base::StartsWith(loc, kLocale));
       break;
     } else {
       for (int i = 0; i < h; ++i, ++y) {
diff --git a/tests/testdata/deflate_src.zip b/tests/testdata/deflate_src.zip
new file mode 100644
index 000000000..bdb2b3216
--- /dev/null
+++ b/tests/testdata/deflate_src.zip
diff --git a/tests/testdata/deflate_tgt.zip b/tests/testdata/deflate_tgt.zip
new file mode 100644
index 000000000..2a21760ec
--- /dev/null
+++ b/tests/testdata/deflate_tgt.zip
diff --git a/tests/unit/dirutil_test.cpp b/tests/unit/dirutil_test.cpp
index 5e2ae4fb5..7f85d13ea 100644
--- a/tests/unit/dirutil_test.cpp
+++ b/tests/unit/dirutil_test.cpp
@@ -26,23 +26,23 @@
 
 TEST(DirUtilTest, create_invalid) {
   // Requesting to create an empty dir is invalid.
-  ASSERT_EQ(-1, dirCreateHierarchy("", 0755, nullptr, false, nullptr));
+  ASSERT_EQ(-1, mkdir_recursively("", 0755, false, nullptr));
   ASSERT_EQ(ENOENT, errno);
 
   // Requesting to strip the name with no slash present.
-  ASSERT_EQ(-1, dirCreateHierarchy("abc", 0755, nullptr, true, nullptr));
+  ASSERT_EQ(-1, mkdir_recursively("abc", 0755, true, nullptr));
   ASSERT_EQ(ENOENT, errno);
 
   // Creating a dir that already exists.
   TemporaryDir td;
-  ASSERT_EQ(0, dirCreateHierarchy(td.path, 0755, nullptr, false, nullptr));
+  ASSERT_EQ(0, mkdir_recursively(td.path, 0755, false, nullptr));
 
   // "///" is a valid dir.
-  ASSERT_EQ(0, dirCreateHierarchy("///", 0755, nullptr, false, nullptr));
+  ASSERT_EQ(0, mkdir_recursively("///", 0755, false, nullptr));
 
   // Request to create a dir, but a file with the same name already exists.
   TemporaryFile tf;
-  ASSERT_EQ(-1, dirCreateHierarchy(tf.path, 0755, nullptr, false, nullptr));
+  ASSERT_EQ(-1, mkdir_recursively(tf.path, 0755, false, nullptr));
   ASSERT_EQ(ENOTDIR, errno);
 }
 
@@ -51,7 +51,7 @@ TEST(DirUtilTest, create_smoke) {
   std::string prefix(td.path);
   std::string path = prefix + "/a/b";
   constexpr mode_t mode = 0755;
-  ASSERT_EQ(0, dirCreateHierarchy(path.c_str(), mode, nullptr, false, nullptr));
+  ASSERT_EQ(0, mkdir_recursively(path, mode, false, nullptr));
 
   // Verify.
   struct stat sb;
@@ -69,7 +69,7 @@ TEST(DirUtilTest, create_strip_filename) {
   TemporaryDir td;
   std::string prefix(td.path);
   std::string path = prefix + "/a/b";
-  ASSERT_EQ(0, dirCreateHierarchy(path.c_str(), 0755, nullptr, true, nullptr));
+  ASSERT_EQ(0, mkdir_recursively(path, 0755, true, nullptr));
 
   // Verify that "../a" exists but not "../a/b".
   struct stat sb;
@@ -83,31 +83,21 @@ TEST(DirUtilTest, create_strip_filename) {
   ASSERT_EQ(0, rmdir((prefix + "/a").c_str()));
 }
 
-TEST(DirUtilTest, create_mode_and_timestamp) {
+TEST(DirUtilTest, create_mode) {
   TemporaryDir td;
   std::string prefix(td.path);
   std::string path = prefix + "/a/b";
-  // Set the timestamp to 8/1/2008.
-  constexpr struct utimbuf timestamp = { 1217592000, 1217592000 };
   constexpr mode_t mode = 0751;
-  ASSERT_EQ(0, dirCreateHierarchy(path.c_str(), mode, &timestamp, false, nullptr));
+  ASSERT_EQ(0, mkdir_recursively(path, mode, false, nullptr));
 
-  // Verify the mode and timestamp for "../a/b".
+  // Verify the mode for "../a/b".
   struct stat sb;
   ASSERT_EQ(0, stat(path.c_str(), &sb)) << strerror(errno);
   ASSERT_TRUE(S_ISDIR(sb.st_mode));
   constexpr mode_t mask = S_IRWXU | S_IRWXG | S_IRWXO;
   ASSERT_EQ(mode, sb.st_mode & mask);
 
-  timespec time;
-  time.tv_sec = 1217592000;
-  time.tv_nsec = 0;
-
-  ASSERT_EQ(time.tv_sec, static_cast<long>(sb.st_atime));
-  ASSERT_EQ(time.tv_sec, static_cast<long>(sb.st_mtime));
-
-  // Verify the mode for "../a". Note that the timestamp for intermediate directories (e.g. "../a")
-  // may not be 'timestamp' according to the current implementation.
+  // Verify the mode for "../a".
   ASSERT_EQ(0, stat((prefix + "/a").c_str(), &sb)) << strerror(errno);
   ASSERT_TRUE(S_ISDIR(sb.st_mode));
   ASSERT_EQ(mode, sb.st_mode & mask);
@@ -116,35 +106,3 @@ TEST(DirUtilTest, create_mode_and_timestamp) {
   ASSERT_EQ(0, rmdir((prefix + "/a/b").c_str()));
   ASSERT_EQ(0, rmdir((prefix + "/a").c_str()));
 }
-
-TEST(DirUtilTest, unlink_invalid) {
-  // File doesn't exist.
-  ASSERT_EQ(-1, dirUnlinkHierarchy("doesntexist"));
-
-  // Nonexistent directory.
-  TemporaryDir td;
-  std::string path(td.path);
-  ASSERT_EQ(-1, dirUnlinkHierarchy((path + "/a").c_str()));
-  ASSERT_EQ(ENOENT, errno);
-}
-
-TEST(DirUtilTest, unlink_smoke) {
-  // Unlink a file.
-  TemporaryFile tf;
-  ASSERT_EQ(0, dirUnlinkHierarchy(tf.path));
-  ASSERT_EQ(-1, access(tf.path, F_OK));
-
-  TemporaryDir td;
-  std::string path(td.path);
-  constexpr mode_t mode = 0700;
-  ASSERT_EQ(0, mkdir((path + "/a").c_str(), mode));
-  ASSERT_EQ(0, mkdir((path + "/a/b").c_str(), mode));
-  ASSERT_EQ(0, mkdir((path + "/a/b/c").c_str(), mode));
-  ASSERT_EQ(0, mkdir((path + "/a/d").c_str(), mode));
-
-  // Remove "../a" recursively.
-  ASSERT_EQ(0, dirUnlinkHierarchy((path + "/a").c_str()));
-
-  // Verify it's gone.
-  ASSERT_EQ(-1, access((path + "/a").c_str(), F_OK));
-}
diff --git a/tests/unit/rangeset_test.cpp b/tests/unit/rangeset_test.cpp
index 3c6d77ef5..7ae193e18 100644
--- a/tests/unit/rangeset_test.cpp
+++ b/tests/unit/rangeset_test.cpp
@@ -17,11 +17,23 @@
 #include <signal.h>
 #include <sys/types.h>
 
+#include <limits>
 #include <vector>
 
 #include <gtest/gtest.h>
 
-#include "updater/rangeset.h"
+#include "otautil/rangeset.h"
+
+TEST(RangeSetTest, ctor) {
+  RangeSet rs(std::vector<Range>{ Range{ 8, 10 }, Range{ 1, 5 } });
+  ASSERT_TRUE(rs);
+
+  RangeSet rs2(std::vector<Range>{});
+  ASSERT_FALSE(rs2);
+
+  RangeSet rs3(std::vector<Range>{ Range{ 8, 10 }, Range{ 5, 1 } });
+  ASSERT_FALSE(rs3);
+}
 
 TEST(RangeSetTest, Parse_smoke) {
   RangeSet rs = RangeSet::Parse("2,1,10");
@@ -37,27 +49,64 @@ TEST(RangeSetTest, Parse_smoke) {
 
   // Leading zeros are fine. But android::base::ParseUint() doesn't like trailing zeros like "10 ".
   ASSERT_EQ(rs, RangeSet::Parse(" 2, 1,   10"));
-  ASSERT_EXIT(RangeSet::Parse("2,1,10 "), ::testing::KilledBySignal(SIGABRT), "");
+  ASSERT_FALSE(RangeSet::Parse("2,1,10 "));
 }
 
 TEST(RangeSetTest, Parse_InvalidCases) {
   // Insufficient number of tokens.
-  ASSERT_EXIT(RangeSet::Parse(""), ::testing::KilledBySignal(SIGABRT), "");
-  ASSERT_EXIT(RangeSet::Parse("2,1"), ::testing::KilledBySignal(SIGABRT), "");
+  ASSERT_FALSE(RangeSet::Parse(""));
+  ASSERT_FALSE(RangeSet::Parse("2,1"));
 
   // The first token (i.e. the number of following tokens) is invalid.
-  ASSERT_EXIT(RangeSet::Parse("a,1,1"), ::testing::KilledBySignal(SIGABRT), "");
-  ASSERT_EXIT(RangeSet::Parse("3,1,1"), ::testing::KilledBySignal(SIGABRT), "");
-  ASSERT_EXIT(RangeSet::Parse("-3,1,1"), ::testing::KilledBySignal(SIGABRT), "");
-  ASSERT_EXIT(RangeSet::Parse("2,1,2,3"), ::testing::KilledBySignal(SIGABRT), "");
+  ASSERT_FALSE(RangeSet::Parse("a,1,1"));
+  ASSERT_FALSE(RangeSet::Parse("3,1,1"));
+  ASSERT_FALSE(RangeSet::Parse("-3,1,1"));
+  ASSERT_FALSE(RangeSet::Parse("2,1,2,3"));
 
   // Invalid tokens.
-  ASSERT_EXIT(RangeSet::Parse("2,1,10a"), ::testing::KilledBySignal(SIGABRT), "");
-  ASSERT_EXIT(RangeSet::Parse("2,,10"), ::testing::KilledBySignal(SIGABRT), "");
+  ASSERT_FALSE(RangeSet::Parse("2,1,10a"));
+  ASSERT_FALSE(RangeSet::Parse("2,,10"));
 
   // Empty or negative range.
-  ASSERT_EXIT(RangeSet::Parse("2,2,2"), ::testing::KilledBySignal(SIGABRT), "");
-  ASSERT_EXIT(RangeSet::Parse("2,2,1"), ::testing::KilledBySignal(SIGABRT), "");
+  ASSERT_FALSE(RangeSet::Parse("2,2,2"));
+  ASSERT_FALSE(RangeSet::Parse("2,2,1"));
+}
+
+TEST(RangeSetTest, Clear) {
+  RangeSet rs = RangeSet::Parse("2,1,6");
+  ASSERT_TRUE(rs);
+  rs.Clear();
+  ASSERT_FALSE(rs);
+
+  // No-op to clear an empty RangeSet.
+  rs.Clear();
+  ASSERT_FALSE(rs);
+}
+
+TEST(RangeSetTest, PushBack) {
+  RangeSet rs;
+  ASSERT_FALSE(rs);
+
+  ASSERT_TRUE(rs.PushBack({ 3, 5 }));
+  ASSERT_EQ(RangeSet::Parse("2,3,5"), rs);
+
+  ASSERT_TRUE(rs.PushBack({ 5, 15 }));
+  ASSERT_EQ(RangeSet::Parse("4,3,5,5,15"), rs);
+  ASSERT_EQ(static_cast<size_t>(2), rs.size());
+  ASSERT_EQ(static_cast<size_t>(12), rs.blocks());
+}
+
+TEST(RangeSetTest, PushBack_InvalidInput) {
+  RangeSet rs;
+  ASSERT_FALSE(rs);
+  ASSERT_FALSE(rs.PushBack({ 5, 3 }));
+  ASSERT_FALSE(rs);
+  ASSERT_FALSE(rs.PushBack({ 15, 15 }));
+  ASSERT_FALSE(rs);
+
+  ASSERT_TRUE(rs.PushBack({ 5, 15 }));
+  ASSERT_FALSE(rs.PushBack({ 5, std::numeric_limits<size_t>::max() - 2 }));
+  ASSERT_EQ(RangeSet::Parse("2,5,15"), rs);
 }
 
 TEST(RangeSetTest, Overlaps) {
@@ -74,6 +123,86 @@ TEST(RangeSetTest, Overlaps) {
   ASSERT_FALSE(RangeSet::Parse("2,5,7").Overlaps(RangeSet::Parse("2,3,5")));
 }
 
+TEST(RangeSetTest, Split) {
+  RangeSet rs1 = RangeSet::Parse("2,1,2");
+  ASSERT_TRUE(rs1);
+  ASSERT_EQ((std::vector<RangeSet>{ RangeSet::Parse("2,1,2") }), rs1.Split(1));
+
+  RangeSet rs2 = RangeSet::Parse("2,5,10");
+  ASSERT_TRUE(rs2);
+  ASSERT_EQ((std::vector<RangeSet>{ RangeSet::Parse("2,5,8"), RangeSet::Parse("2,8,10") }),
+            rs2.Split(2));
+
+  RangeSet rs3 = RangeSet::Parse("4,0,1,5,10");
+  ASSERT_TRUE(rs3);
+  ASSERT_EQ((std::vector<RangeSet>{ RangeSet::Parse("4,0,1,5,7"), RangeSet::Parse("2,7,10") }),
+            rs3.Split(2));
+
+  RangeSet rs4 = RangeSet::Parse("6,1,3,3,4,4,5");
+  ASSERT_TRUE(rs4);
+  ASSERT_EQ((std::vector<RangeSet>{ RangeSet::Parse("2,1,3"), RangeSet::Parse("2,3,4"),
+                                    RangeSet::Parse("2,4,5") }),
+            rs4.Split(3));
+
+  RangeSet rs5 = RangeSet::Parse("2,0,10");
+  ASSERT_TRUE(rs5);
+  ASSERT_EQ((std::vector<RangeSet>{ RangeSet::Parse("2,0,3"), RangeSet::Parse("2,3,6"),
+                                    RangeSet::Parse("2,6,8"), RangeSet::Parse("2,8,10") }),
+            rs5.Split(4));
+
+  RangeSet rs6 = RangeSet::Parse(
+      "20,0,268,269,271,286,447,8350,32770,33022,98306,98558,163842,164094,196609,204800,229378,"
+      "229630,294914,295166,457564");
+  ASSERT_TRUE(rs6);
+  size_t rs6_blocks = rs6.blocks();
+  auto splits = rs6.Split(4);
+  ASSERT_EQ(
+      (std::vector<RangeSet>{
+          RangeSet::Parse("12,0,268,269,271,286,447,8350,32770,33022,98306,98558,118472"),
+          RangeSet::Parse("8,118472,163842,164094,196609,204800,229378,229630,237216"),
+          RangeSet::Parse("4,237216,294914,295166,347516"), RangeSet::Parse("2,347516,457564") }),
+      splits);
+  size_t sum = 0;
+  for (const auto& element : splits) {
+    sum += element.blocks();
+  }
+  ASSERT_EQ(rs6_blocks, sum);
+}
+
+TEST(RangeSetTest, Split_EdgeCases) {
+  // Empty RangeSet.
+  RangeSet rs1;
+  ASSERT_FALSE(rs1);
+  ASSERT_EQ((std::vector<RangeSet>{}), rs1.Split(2));
+  ASSERT_FALSE(rs1);
+
+  // Zero group.
+  RangeSet rs2 = RangeSet::Parse("2,1,5");
+  ASSERT_TRUE(rs2);
+  ASSERT_EQ((std::vector<RangeSet>{}), rs2.Split(0));
+
+  // The number of blocks equals to the number of groups.
+  RangeSet rs3 = RangeSet::Parse("2,1,5");
+  ASSERT_TRUE(rs3);
+  ASSERT_EQ((std::vector<RangeSet>{ RangeSet::Parse("2,1,2"), RangeSet::Parse("2,2,3"),
+                                    RangeSet::Parse("2,3,4"), RangeSet::Parse("2,4,5") }),
+            rs3.Split(4));
+
+  // Less blocks than the number of groups.
+  RangeSet rs4 = RangeSet::Parse("2,1,5");
+  ASSERT_TRUE(rs4);
+  ASSERT_EQ((std::vector<RangeSet>{ RangeSet::Parse("2,1,2"), RangeSet::Parse("2,2,3"),
+                                    RangeSet::Parse("2,3,4"), RangeSet::Parse("2,4,5") }),
+            rs4.Split(8));
+
+  // Less blocks than the number of groups.
+  RangeSet rs5 = RangeSet::Parse("2,0,3");
+  ASSERT_TRUE(rs5);
+  ASSERT_EQ((std::vector<RangeSet>{ RangeSet::Parse("2,0,1"), RangeSet::Parse("2,1,2"),
+                                    RangeSet::Parse("2,2,3") }),
+            rs5.Split(4));
+}
+
 TEST(RangeSetTest, GetBlockNumber) {
   RangeSet rs = RangeSet::Parse("2,1,10");
   ASSERT_EQ(static_cast<size_t>(1), rs.GetBlockNumber(0));
@@ -90,7 +219,7 @@ TEST(RangeSetTest, equality) {
   ASSERT_NE(RangeSet::Parse("2,1,6"), RangeSet::Parse("2,1,7"));
   ASSERT_NE(RangeSet::Parse("2,1,6"), RangeSet::Parse("2,2,7"));
 
-  // The orders of Range's matter. "4,1,5,8,10" != "4,8,10,1,5".
+  // The orders of Range's matter, e.g. "4,1,5,8,10" != "4,8,10,1,5".
   ASSERT_NE(RangeSet::Parse("4,1,5,8,10"), RangeSet::Parse("4,8,10,1,5"));
 }
 
@@ -110,3 +239,51 @@ TEST(RangeSetTest, iterators) {
   }
   ASSERT_EQ((std::vector<Range>{ Range{ 8, 10 }, Range{ 1, 5 } }), ranges);
 }
+
+TEST(RangeSetTest, ToString) {
+  ASSERT_EQ("", RangeSet::Parse("").ToString());
+  ASSERT_EQ("2,1,6", RangeSet::Parse("2,1,6").ToString());
+  ASSERT_EQ("4,1,5,8,10", RangeSet::Parse("4,1,5,8,10").ToString());
+  ASSERT_EQ("6,1,3,4,6,15,22", RangeSet::Parse("6,1,3,4,6,15,22").ToString());
+}
+
+TEST(SortedRangeSetTest, Insert) {
+  SortedRangeSet rs({ { 2, 3 }, { 4, 6 }, { 8, 14 } });
+  rs.Insert({ 1, 2 });
+  ASSERT_EQ(SortedRangeSet({ { 1, 3 }, { 4, 6 }, { 8, 14 } }), rs);
+  ASSERT_EQ(static_cast<size_t>(10), rs.blocks());
+  rs.Insert({ 3, 5 });
+  ASSERT_EQ(SortedRangeSet({ { 1, 6 }, { 8, 14 } }), rs);
+  ASSERT_EQ(static_cast<size_t>(11), rs.blocks());
+
+  SortedRangeSet r1({ { 20, 22 }, { 15, 18 } });
+  rs.Insert(r1);
+  ASSERT_EQ(SortedRangeSet({ { 1, 6 }, { 8, 14 }, { 15, 18 }, { 20, 22 } }), rs);
+  ASSERT_EQ(static_cast<size_t>(16), rs.blocks());
+
+  SortedRangeSet r2({ { 2, 7 }, { 15, 21 }, { 20, 25 } });
+  rs.Insert(r2);
+  ASSERT_EQ(SortedRangeSet({ { 1, 7 }, { 8, 14 }, { 15, 25 } }), rs);
+  ASSERT_EQ(static_cast<size_t>(22), rs.blocks());
+}
+
+TEST(SortedRangeSetTest, file_range) {
+  SortedRangeSet rs;
+  rs.Insert(4096, 4096);
+  ASSERT_EQ(SortedRangeSet({ { 1, 2 } }), rs);
+  // insert block 2-9
+  rs.Insert(4096 * 3 - 1, 4096 * 7);
+  ASSERT_EQ(SortedRangeSet({ { 1, 10 } }), rs);
+  // insert block 15-19
+  rs.Insert(4096 * 15 + 1, 4096 * 4);
+  ASSERT_EQ(SortedRangeSet({ { 1, 10 }, { 15, 20 } }), rs);
+
+  // rs overlaps block 2-2
+  ASSERT_TRUE(rs.Overlaps(4096 * 2 - 1, 10));
+  ASSERT_FALSE(rs.Overlaps(4096 * 10, 4096 * 5));
+
+  ASSERT_EQ(static_cast<size_t>(10), rs.GetOffsetInRangeSet(4106));
+  ASSERT_EQ(static_cast<size_t>(40970), rs.GetOffsetInRangeSet(4096 * 16 + 10));
+  // block#10 not in range.
+  ASSERT_EXIT(rs.GetOffsetInRangeSet(40970), ::testing::KilledBySignal(SIGABRT), "");
+}