diff --git a/Android.mk b/Android.mk
index 257f776..6aa91ea 100644
--- a/Android.mk
+++ b/Android.mk
@@ -97,61 +97,7 @@
 
 include $(BUILD_STATIC_LIBRARY)
 
-# librecovery (static library)
-# ===============================
-include $(CLEAR_VARS)
-
-LOCAL_SRC_FILES := \
-    install.cpp
-
-LOCAL_CFLAGS := $(recovery_common_cflags)
-
-ifeq ($(AB_OTA_UPDATER),true)
-    LOCAL_CFLAGS += -DAB_OTA_UPDATER=1
-endif
-
-LOCAL_MODULE := librecovery
-
-LOCAL_STATIC_LIBRARIES := \
-    libminui \
-    libotautil \
-    libvintf_recovery \
-    libcrypto_utils \
-    libcrypto \
-    libbase \
-    libziparchive \
-
-include $(BUILD_STATIC_LIBRARY)
-
-# recovery (static executable)
-# ===============================
-include $(CLEAR_VARS)
-
-LOCAL_SRC_FILES := \
-    adb_install.cpp \
-    fuse_sdcard_provider.cpp \
-    logging.cpp \
-    recovery.cpp \
-    recovery_main.cpp \
-    roots.cpp \
-
-LOCAL_MODULE := recovery
-
-LOCAL_FORCE_STATIC_EXECUTABLE := true
-
-LOCAL_MODULE_PATH := $(TARGET_RECOVERY_ROOT_OUT)/sbin
-
-# Cannot link with LLD: undefined symbol: UsbNoPermissionsLongHelpText
-# http://b/77543887, lld does not handle -Wl,--gc-sections as well as ld.
-LOCAL_USE_CLANG_LLD := false
-
-LOCAL_CFLAGS := $(recovery_common_cflags)
-
-LOCAL_C_INCLUDES += \
-    system/vold \
-
-LOCAL_STATIC_LIBRARIES := \
-    librecovery \
+librecovery_static_libraries := \
     $(TARGET_RECOVERY_UI_LIB) \
     libbootloader_message \
     libfusesideload \
@@ -174,12 +120,63 @@
     libtinyxml2 \
     libziparchive \
     libbase \
-    libcutils \
     libutils \
+    libcutils \
     liblog \
     libselinux \
     libz \
 
+# librecovery (static library)
+# ===============================
+include $(CLEAR_VARS)
+
+LOCAL_SRC_FILES := \
+    adb_install.cpp \
+    fuse_sdcard_provider.cpp \
+    install.cpp \
+    recovery.cpp \
+    roots.cpp \
+
+LOCAL_C_INCLUDES := \
+    system/vold \
+
+LOCAL_CFLAGS := $(recovery_common_cflags)
+
+ifeq ($(AB_OTA_UPDATER),true)
+    LOCAL_CFLAGS += -DAB_OTA_UPDATER=1
+endif
+
+LOCAL_MODULE := librecovery
+
+LOCAL_STATIC_LIBRARIES := \
+    $(librecovery_static_libraries)
+
+include $(BUILD_STATIC_LIBRARY)
+
+# recovery (static executable)
+# ===============================
+include $(CLEAR_VARS)
+
+LOCAL_SRC_FILES := \
+    logging.cpp \
+    recovery_main.cpp \
+
+LOCAL_MODULE := recovery
+
+LOCAL_FORCE_STATIC_EXECUTABLE := true
+
+LOCAL_MODULE_PATH := $(TARGET_RECOVERY_ROOT_OUT)/sbin
+
+# Cannot link with LLD: undefined symbol: UsbNoPermissionsLongHelpText
+# http://b/77543887, lld does not handle -Wl,--gc-sections as well as ld.
+LOCAL_USE_CLANG_LLD := false
+
+LOCAL_CFLAGS := $(recovery_common_cflags)
+
+LOCAL_STATIC_LIBRARIES := \
+    librecovery \
+    $(librecovery_static_libraries)
+
 LOCAL_HAL_STATIC_LIBRARIES := libhealthd
 
 LOCAL_REQUIRED_MODULES := \
diff --git a/applypatch/imgdiff.cpp b/applypatch/imgdiff.cpp
index 674cc2b..415d95f 100644
--- a/applypatch/imgdiff.cpp
+++ b/applypatch/imgdiff.cpp
@@ -462,12 +462,12 @@
       target_len_(tgt.GetRawDataLength()),
       target_uncompressed_len_(tgt.DataLengthForPatch()),
       target_compress_level_(tgt.GetCompressLevel()),
-      data_(tgt.DataForPatch(), tgt.DataForPatch() + tgt.DataLengthForPatch()) {}
+      data_(tgt.GetRawData(), tgt.GetRawData() + tgt.GetRawDataLength()) {}
 
 // Return true if raw data is smaller than the patch size.
 bool PatchChunk::RawDataIsSmaller(const ImageChunk& tgt, size_t patch_size) {
   size_t target_len = tgt.GetRawDataLength();
-  return (tgt.GetType() == CHUNK_NORMAL && (target_len <= 160 || target_len < patch_size));
+  return target_len < patch_size || (tgt.GetType() == CHUNK_NORMAL && target_len <= 160);
 }
 
 void PatchChunk::UpdateSourceOffset(const SortedRangeSet& src_range) {
diff --git a/applypatch/imgpatch.cpp b/applypatch/imgpatch.cpp
index c4c2707..2f8f485 100644
--- a/applypatch/imgpatch.cpp
+++ b/applypatch/imgpatch.cpp
@@ -54,6 +54,7 @@
                                             const Value& patch, size_t patch_offset,
                                             const char* deflate_header, SinkFn sink) {
   size_t expected_target_length = static_cast<size_t>(Read8(deflate_header + 32));
+  CHECK_GT(expected_target_length, 0);
   int level = Read4(deflate_header + 40);
   int method = Read4(deflate_header + 44);
   int window_bits = Read4(deflate_header + 48);
diff --git a/applypatch/include/applypatch/imgdiff_image.h b/applypatch/include/applypatch/imgdiff_image.h
index 0848072..6716051 100644
--- a/applypatch/include/applypatch/imgdiff_image.h
+++ b/applypatch/include/applypatch/imgdiff_image.h
@@ -44,6 +44,8 @@
   int GetType() const {
     return type_;
   }
+
+  const uint8_t* GetRawData() const;
   size_t GetRawDataLength() const {
     return raw_data_len_;
   }
@@ -99,7 +101,6 @@
                         bsdiff::SuffixArrayIndexInterface** bsdiff_cache);
 
  private:
-  const uint8_t* GetRawData() const;
   bool TryReconstruction(int level);
 
   int type_;                                    // CHUNK_NORMAL, CHUNK_DEFLATE, CHUNK_RAW
diff --git a/bootloader_message/Android.bp b/bootloader_message/Android.bp
index c81c67b..ab23733 100644
--- a/bootloader_message/Android.bp
+++ b/bootloader_message/Android.bp
@@ -16,6 +16,7 @@
 
 cc_library_static {
     name: "libbootloader_message",
+    recovery_available: true,
     srcs: ["bootloader_message.cpp"],
     cflags: [
         "-Wall",
diff --git a/recovery.cpp b/recovery.cpp
index b1a2900..56b2567 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -1049,6 +1049,7 @@
         continue;
     }
   }
+  optind = 1;
 
   printf("stage is [%s]\n", stage.c_str());
   printf("reason is [%s]\n", reason);
@@ -1062,6 +1063,11 @@
     ui->SetStage(st_cur, st_max);
   }
 
+  std::vector<std::string> title_lines =
+      android::base::Split(android::base::GetProperty("ro.bootimage.build.fingerprint", ""), ":");
+  title_lines.insert(std::begin(title_lines), "Android Recovery");
+  ui->SetTitle(title_lines);
+
   device->StartRecovery();
 
   printf("Command:");
diff --git a/recovery_main.cpp b/recovery_main.cpp
index e21c782..5e82c6c 100644
--- a/recovery_main.cpp
+++ b/recovery_main.cpp
@@ -317,6 +317,7 @@
       }
     }
   }
+  optind = 1;
 
   if (locale.empty()) {
     if (has_cache) {
diff --git a/screen_ui.cpp b/screen_ui.cpp
index fd7a1be..f1b3878 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -496,6 +496,10 @@
   return offset;
 }
 
+void ScreenRecoveryUI::SetTitle(const std::vector<std::string>& lines) {
+  title_lines_ = lines;
+}
+
 // Redraws everything on the screen. Does not flip pages. Should only be called with updateMutex
 // locked.
 void ScreenRecoveryUI::draw_screen_locked() {
@@ -529,11 +533,9 @@
     int x = kMarginWidth + kMenuIndent;
 
     SetColor(INFO);
-    y += DrawTextLine(x, y, "Android Recovery", true);
-    std::string recovery_fingerprint =
-        android::base::GetProperty("ro.bootimage.build.fingerprint", "");
-    for (const auto& chunk : android::base::Split(recovery_fingerprint, ":")) {
-      y += DrawTextLine(x, y, chunk, false);
+
+    for (size_t i = 0; i < title_lines_.size(); i++) {
+      y += DrawTextLine(x, y, title_lines_[i], i == 0);
     }
 
     y += DrawTextLines(x, y, help_message);
diff --git a/screen_ui.h b/screen_ui.h
index 2d6b621..c90a2cd 100644
--- a/screen_ui.h
+++ b/screen_ui.h
@@ -141,6 +141,7 @@
   size_t ShowMenu(const std::vector<std::string>& headers, const std::vector<std::string>& items,
                   size_t initial_selection, bool menu_only,
                   const std::function<int(int, bool)>& key_handler) override;
+  void SetTitle(const std::vector<std::string>& lines) override;
 
   void KeyLongPress(int) override;
 
@@ -266,6 +267,8 @@
   bool show_text;
   bool show_text_ever;  // has show_text ever been true?
 
+  std::vector<std::string> title_lines_;
+
   bool scrollable_menu_;
   std::unique_ptr<Menu> menu_;
 
diff --git a/stub_ui.h b/stub_ui.h
index 67c338e..a3cf12b 100644
--- a/stub_ui.h
+++ b/stub_ui.h
@@ -67,6 +67,8 @@
                   const std::function<int(int, bool)>& /* key_handler */) override {
     return initial_selection;
   }
+
+  void SetTitle(const std::vector<std::string>& /* lines */) override {}
 };
 
 #endif  // RECOVERY_STUB_UI_H
diff --git a/tests/Android.mk b/tests/Android.mk
index cdc5b52..853ca27 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -174,8 +174,8 @@
     libtinyxml2 \
     libziparchive \
     libbase \
-    libcutils \
     libutils \
+    libcutils \
     liblog \
     libselinux \
     libz \
diff --git a/tests/component/imgdiff_test.cpp b/tests/component/imgdiff_test.cpp
index 6c23def..cb4868a 100644
--- a/tests/component/imgdiff_test.cpp
+++ b/tests/component/imgdiff_test.cpp
@@ -197,12 +197,17 @@
 }
 
 TEST(ImgdiffTest, zip_mode_smoke_compressed) {
+  // Generate 1 block of random data.
+  std::string random_data;
+  random_data.reserve(4096);
+  generate_n(back_inserter(random_data), 4096, []() { return rand() % 256; });
+
   // Construct src and tgt zip files.
   TemporaryFile src_file;
   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");
+  const std::string src_content = random_data;
   ASSERT_EQ(0, src_writer.WriteBytes(src_content.data(), src_content.size()));
   ASSERT_EQ(0, src_writer.FinishEntry());
   ASSERT_EQ(0, src_writer.Finish());
@@ -212,7 +217,7 @@
   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");
+  const std::string tgt_content = random_data + "extra contents";
   ASSERT_EQ(0, tgt_writer.WriteBytes(tgt_content.data(), tgt_content.size()));
   ASSERT_EQ(0, tgt_writer.FinishEntry());
   ASSERT_EQ(0, tgt_writer.Finish());
@@ -245,13 +250,57 @@
   verify_patched_image(src, patch, tgt);
 }
 
+TEST(ImgdiffTest, zip_mode_empty_target) {
+  TemporaryFile src_file;
+  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";
+  ASSERT_EQ(0, src_writer.WriteBytes(src_content.data(), src_content.size()));
+  ASSERT_EQ(0, src_writer.FinishEntry());
+  ASSERT_EQ(0, src_writer.Finish());
+  ASSERT_EQ(0, fclose(src_file_ptr));
+
+  // Construct a empty entry in the target zip.
+  TemporaryFile tgt_file;
+  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;
+  ASSERT_EQ(0, tgt_writer.WriteBytes(tgt_content.data(), tgt_content.size()));
+  ASSERT_EQ(0, tgt_writer.FinishEntry());
+  ASSERT_EQ(0, tgt_writer.Finish());
+
+  // Compute patch.
+  TemporaryFile patch_file;
+  std::vector<const char*> args = {
+    "imgdiff", "-z", src_file.path, tgt_file.path, patch_file.path,
+  };
+  ASSERT_EQ(0, imgdiff(args.size(), args.data()));
+
+  // Verify.
+  std::string tgt;
+  ASSERT_TRUE(android::base::ReadFileToString(tgt_file.path, &tgt));
+  std::string src;
+  ASSERT_TRUE(android::base::ReadFileToString(src_file.path, &src));
+  std::string patch;
+  ASSERT_TRUE(android::base::ReadFileToString(patch_file.path, &patch));
+
+  verify_patched_image(src, patch, tgt);
+}
+
 TEST(ImgdiffTest, zip_mode_smoke_trailer_zeros) {
+  // Generate 1 block of random data.
+  std::string random_data;
+  random_data.reserve(4096);
+  generate_n(back_inserter(random_data), 4096, []() { return rand() % 256; });
+
   // Construct src and tgt zip files.
   TemporaryFile src_file;
   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");
+  const std::string src_content = random_data;
   ASSERT_EQ(0, src_writer.WriteBytes(src_content.data(), src_content.size()));
   ASSERT_EQ(0, src_writer.FinishEntry());
   ASSERT_EQ(0, src_writer.Finish());
@@ -261,7 +310,7 @@
   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");
+  const std::string tgt_content = random_data + "abcdefg";
   ASSERT_EQ(0, tgt_writer.WriteBytes(tgt_content.data(), tgt_content.size()));
   ASSERT_EQ(0, tgt_writer.FinishEntry());
   ASSERT_EQ(0, tgt_writer.Finish());
@@ -298,23 +347,19 @@
 }
 
 TEST(ImgdiffTest, image_mode_simple) {
-  // src: "abcdefgh" + gzipped "xyz" (echo -n "xyz" | gzip -f | hd).
-  const std::vector<char> src_data = { 'a',    'b',    'c',    'd',    'e',    'f',    'g',
-                                       'h',    '\x1f', '\x8b', '\x08', '\x00', '\xc4', '\x1e',
-                                       '\x53', '\x58', '\x00', '\x03', '\xab', '\xa8', '\xac',
-                                       '\x02', '\x00', '\x67', '\xba', '\x8e', '\xeb', '\x03',
-                                       '\x00', '\x00', '\x00' };
-  const std::string src(src_data.cbegin(), src_data.cend());
+  std::string gzipped_source_path = from_testdata_base("gzipped_source");
+  std::string gzipped_source;
+  ASSERT_TRUE(android::base::ReadFileToString(gzipped_source_path, &gzipped_source));
+
+  const std::string src = "abcdefg" + gzipped_source;
   TemporaryFile src_file;
   ASSERT_TRUE(android::base::WriteStringToFile(src, src_file.path));
 
-  // tgt: "abcdefgxyz" + gzipped "xxyyzz".
-  const std::vector<char> tgt_data = {
-    'a',    'b',    'c',    'd',    'e',    'f',    'g',    'x',    'y',    'z',    '\x1f', '\x8b',
-    '\x08', '\x00', '\x62', '\x1f', '\x53', '\x58', '\x00', '\x03', '\xab', '\xa8', '\xa8', '\xac',
-    '\xac', '\xaa', '\x02', '\x00', '\x96', '\x30', '\x06', '\xb7', '\x06', '\x00', '\x00', '\x00'
-  };
-  const std::string tgt(tgt_data.cbegin(), tgt_data.cend());
+  std::string gzipped_target_path = from_testdata_base("gzipped_target");
+  std::string gzipped_target;
+  ASSERT_TRUE(android::base::ReadFileToString(gzipped_target_path, &gzipped_target));
+  const std::string tgt = "abcdefgxyz" + gzipped_target;
+
   TemporaryFile tgt_file;
   ASSERT_TRUE(android::base::WriteStringToFile(tgt, tgt_file.path));
 
@@ -404,23 +449,21 @@
 }
 
 TEST(ImgdiffTest, image_mode_merge_chunks) {
-  // src: "abcdefgh" + gzipped "xyz" (echo -n "xyz" | gzip -f | hd).
-  const std::vector<char> src_data = { 'a',    'b',    'c',    'd',    'e',    'f',    'g',
-                                       'h',    '\x1f', '\x8b', '\x08', '\x00', '\xc4', '\x1e',
-                                       '\x53', '\x58', '\x00', '\x03', '\xab', '\xa8', '\xac',
-                                       '\x02', '\x00', '\x67', '\xba', '\x8e', '\xeb', '\x03',
-                                       '\x00', '\x00', '\x00' };
-  const std::string src(src_data.cbegin(), src_data.cend());
+  // src: "abcdefg" + gzipped_source.
+  std::string gzipped_source_path = from_testdata_base("gzipped_source");
+  std::string gzipped_source;
+  ASSERT_TRUE(android::base::ReadFileToString(gzipped_source_path, &gzipped_source));
+
+  const std::string src = "abcdefg" + gzipped_source;
   TemporaryFile src_file;
   ASSERT_TRUE(android::base::WriteStringToFile(src, src_file.path));
 
-  // tgt: gzipped "xyz" + "abcdefgh".
-  const std::vector<char> tgt_data = {
-    '\x1f', '\x8b', '\x08', '\x00', '\x62', '\x1f', '\x53', '\x58', '\x00', '\x03', '\xab', '\xa8',
-    '\xa8', '\xac', '\xac', '\xaa', '\x02', '\x00', '\x96', '\x30', '\x06', '\xb7', '\x06', '\x00',
-    '\x00', '\x00', 'a',    'b',    'c',    'd',    'e',    'f',    'g',    'x',    'y',    'z'
-  };
-  const std::string tgt(tgt_data.cbegin(), tgt_data.cend());
+  // tgt: gzipped_target + "abcdefgxyz".
+  std::string gzipped_target_path = from_testdata_base("gzipped_target");
+  std::string gzipped_target;
+  ASSERT_TRUE(android::base::ReadFileToString(gzipped_target_path, &gzipped_target));
+
+  const std::string tgt = gzipped_target + "abcdefgxyz";
   TemporaryFile tgt_file;
   ASSERT_TRUE(android::base::WriteStringToFile(tgt, tgt_file.path));
 
diff --git a/tests/component/updater_test.cpp b/tests/component/updater_test.cpp
index 48363a6..de8fafd 100644
--- a/tests/component/updater_test.cpp
+++ b/tests/component/updater_test.cpp
@@ -27,6 +27,8 @@
 #include <vector>
 
 #include <android-base/file.h>
+#include <android-base/logging.h>
+#include <android-base/parseint.h>
 #include <android-base/properties.h>
 #include <android-base/stringprintf.h>
 #include <android-base/strings.h>
@@ -48,7 +50,11 @@
 #include "updater/install.h"
 #include "updater/updater.h"
 
-struct selabel_handle *sehandle = nullptr;
+using PackageEntries = std::unordered_map<std::string, std::string>;
+
+static constexpr size_t kTransferListHeaderLines = 4;
+
+struct selabel_handle* sehandle = nullptr;
 
 static void expect(const char* expected, const char* expr_str, CauseCode cause_code,
                    UpdaterInfo* info = nullptr) {
@@ -76,12 +82,12 @@
   ASSERT_EQ(cause_code, state.cause_code);
 }
 
-static void BuildUpdatePackage(const std::unordered_map<std::string, std::string>& entries,
-                               int fd) {
+static void BuildUpdatePackage(const PackageEntries& entries, int fd) {
   FILE* zip_file_ptr = fdopen(fd, "wb");
   ZipWriter zip_writer(zip_file_ptr);
 
   for (const auto& entry : entries) {
+    // All the entries are written as STORED.
     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()));
@@ -93,6 +99,37 @@
   ASSERT_EQ(0, fclose(zip_file_ptr));
 }
 
+static void RunBlockImageUpdate(bool is_verify, const PackageEntries& entries,
+                                const std::string& image_file, const std::string& result) {
+  CHECK(entries.find("transfer_list") != entries.end());
+
+  // 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 new_data = entries.find("new_data.br") != entries.end() ? "new_data.br" : "new_data";
+  std::string script = is_verify ? "block_image_verify" : "block_image_update";
+  script += R"((")" + image_file + R"(", package_extract_file("transfer_list"), ")" + new_data +
+            R"(", "patch_data"))";
+  expect(result.c_str(), script.c_str(), kNoCause, &updater_info);
+
+  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
+  CloseArchive(handle);
+}
+
 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);
@@ -101,19 +138,39 @@
 
 class UpdaterTest : public ::testing::Test {
  protected:
-  virtual void SetUp() override {
+  void SetUp() override {
     RegisterBuiltins();
     RegisterInstallFunctions();
     RegisterBlockImageFunctions();
 
+    // Each test is run in a separate process (isolated mode). Shared temporary files won't cause
+    // conflicts.
     Paths::Get().set_cache_temp_source(temp_saved_source_.path);
     Paths::Get().set_last_command_file(temp_last_command_.path);
     Paths::Get().set_stash_directory_base(temp_stash_base_.path);
+
+    last_command_file_ = temp_last_command_.path;
+    image_file_ = image_temp_file_.path;
+  }
+
+  void TearDown() override {
+    // Clean up the last_command_file if any.
+    ASSERT_TRUE(android::base::RemoveFileIfExists(last_command_file_));
+
+    // Clear partition updated marker if any.
+    std::string updated_marker{ temp_stash_base_.path };
+    updated_marker += "/" + get_sha1(image_temp_file_.path) + ".UPDATED";
+    ASSERT_TRUE(android::base::RemoveFileIfExists(updated_marker));
   }
 
   TemporaryFile temp_saved_source_;
-  TemporaryFile temp_last_command_;
   TemporaryDir temp_stash_base_;
+  std::string last_command_file_;
+  std::string image_file_;
+
+ private:
+  TemporaryFile temp_last_command_;
+  TemporaryFile image_temp_file_;
 };
 
 TEST_F(UpdaterTest, getprop) {
@@ -453,16 +510,18 @@
 
   // 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));
+  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));
 
   // 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 = {
+  std::vector<std::string> transfer_list{
+    // clang-format off
     "4",
     "2",
     "0",
@@ -471,183 +530,108 @@
     android::base::StringPrintf("bsdiff 0 %zu %s %s 2,0,2 2 - %s:2,0,2", patch_content.size(),
                                 src_hash.c_str(), tgt_hash.c_str(), src_hash.c_str()),
     "free " + src_hash,
+    // clang-format on
   };
 
-  std::unordered_map<std::string, std::string> entries = {
+  PackageEntries entries{
     { "new_data", "" },
     { "patch_data", patch_content },
     { "transfer_list", android::base::Join(transfer_list, '\n') },
   };
 
-  // Build the update package.
-  TemporaryFile zip_file;
-  BuildUpdatePackage(entries, zip_file.release());
+  ASSERT_TRUE(android::base::WriteStringToFile(src_content, image_file_));
 
-  MemMapping map;
-  ASSERT_TRUE(map.MapFile(zip_file.path));
-  ZipArchiveHandle handle;
-  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));
+  RunBlockImageUpdate(false, entries, image_file_, "t");
 
-  // 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;
-
-  // 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) +
-      R"(", package_extract_file("transfer_list"), "new_data", "patch_data"))";
-  expect("t", script.c_str(), kNoCause, &updater_info);
   // The update_file should be patched correctly.
   std::string updated_content;
-  ASSERT_TRUE(android::base::ReadFileToString(update_file.path, &updated_content));
-  ASSERT_EQ(tgt_hash, get_sha1(updated_content));
-
-  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
-  CloseArchive(handle);
+  ASSERT_TRUE(android::base::ReadFileToString(image_file_, &updated_content));
+  ASSERT_EQ(tgt_content, updated_content);
 }
 
 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",
+  std::vector<std::string> transfer_list{
+    // clang-format off
+    "4",
+    "2",
+    "0",
+    "2",
+    "stash " + src_hash + " 2,0,2",
+    "free " + src_hash,
+    "fail",
+    // clang-format on
   };
 
   // Add a new data of 10 bytes to test the deadlock.
-  std::unordered_map<std::string, std::string> entries = {
+  PackageEntries 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());
+  ASSERT_TRUE(android::base::WriteStringToFile(src_content, image_file_));
 
-  MemMapping map;
-  ASSERT_TRUE(map.MapFile(zip_file.path));
-  ZipArchiveHandle handle;
-  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));
+  RunBlockImageUpdate(false, entries, image_file_, "");
 
-  // 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 name_digest = get_sha1(image_file_);
   std::string stash_base = std::string(temp_stash_base_.path) + "/" + name_digest;
   ASSERT_EQ(0, access(stash_base.c_str(), F_OK));
+  // Expect the stashed blocks to be freed.
   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_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;
-  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;
-  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 = {
+  std::vector<std::string> transfer_list{
+    // clang-format off
     "4",
     "1",
     "0",
     "0",
     "new 2,0,1",
+    // clang-format on
   };
 
-  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') },
+  // Write 4096 + 100 bytes of new data.
+  PackageEntries entries{
+    { "new_data", std::string(4196, 0) },
     { "patch_data", "" },
     { "transfer_list", android::base::Join(transfer_list, '\n') },
   };
 
-  TemporaryFile zip_file;
-  BuildUpdatePackage(entries, zip_file.release());
+  RunBlockImageUpdate(false, entries, image_file_, "t");
+}
 
-  MemMapping map;
-  ASSERT_TRUE(map.MapFile(zip_file.path));
-  ZipArchiveHandle handle;
-  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));
+TEST_F(UpdaterTest, new_data_short_write) {
+  std::vector<std::string> transfer_list{
+    // clang-format off
+    "4",
+    "1",
+    "0",
+    "0",
+    "new 2,0,1",
+    // clang-format on
+  };
 
-  // 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;
+  PackageEntries entries{
+    { "patch_data", "" },
+    { "transfer_list", android::base::Join(transfer_list, '\n') },
+  };
 
   // Updater should report the failure gracefully rather than stuck in deadlock.
-  TemporaryFile update_file;
-  std::string script_empty_data = "block_image_update(\"" + std::string(update_file.path) +
-      R"(", package_extract_file("transfer_list"), "empty_new_data", "patch_data"))";
-  expect("", script_empty_data.c_str(), kNoCause, &updater_info);
+  entries["new_data"] = "";
+  RunBlockImageUpdate(false, entries, image_file_, "");
 
-  std::string script_short_data = "block_image_update(\"" + std::string(update_file.path) +
-      R"(", package_extract_file("transfer_list"), "short_new_data", "patch_data"))";
-  expect("", script_short_data.c_str(), kNoCause, &updater_info);
+  entries["new_data"] = std::string(10, 'a');
+  RunBlockImageUpdate(false, entries, image_file_, "");
 
   // Expect to write 1 block of new data successfully.
-  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);
+  entries["new_data"] = std::string(4096, 'a');
+  RunBlockImageUpdate(false, entries, image_file_, "t");
 }
 
 TEST_F(UpdaterTest, brotli_new_data) {
@@ -680,55 +664,30 @@
     "new 2,99,100",
   };
 
-  std::unordered_map<std::string, std::string> entries = {
-    { "new.dat.br", std::move(encoded_data) },
+  PackageEntries entries{
+    { "new_data.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));
-  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(), "wb");
-  updater_info.package_zip_addr = map.addr;
-  updater_info.package_zip_len = map.length;
-
-  // Check if we can decompress the new data correctly.
-  TemporaryFile update_file;
-  std::string script_new_data =
-      "block_image_update(\"" + std::string(update_file.path) +
-      R"(", package_extract_file("transfer_list"), "new.dat.br", "patch_data"))";
-  expect("t", script_new_data.c_str(), kNoCause, &updater_info);
+  RunBlockImageUpdate(false, entries, image_file_, "t");
 
   std::string updated_content;
-  ASSERT_TRUE(android::base::ReadFileToString(update_file.path, &updated_content));
+  ASSERT_TRUE(android::base::ReadFileToString(image_file_, &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 = Paths::Get().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(4096, '1');
+  std::string block2(4096, '2');
+  std::string block3(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 = {
+  std::vector<std::string> transfer_list_fail{
+    // clang-format off
     "4",
     "2",
     "0",
@@ -737,10 +696,12 @@
     "move " + block1_hash + " 2,1,2 1 2,0,1",
     "stash " + block3_hash + " 2,2,3",
     "fail",
+    // clang-format on
   };
 
   // Mimic a resumed update with the same transfer commands.
-  std::vector<std::string> transfer_list_continue = {
+  std::vector<std::string> transfer_list_continue{
+    // clang-format off
     "4",
     "2",
     "0",
@@ -749,127 +710,86 @@
     "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",
+    // clang-format on
   };
 
-  std::unordered_map<std::string, std::string> entries = {
+  ASSERT_TRUE(android::base::WriteStringToFile(block1 + block2 + block3, image_file_));
+
+  PackageEntries 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') },
+    { "transfer_list", android::base::Join(transfer_list_fail, '\n') },
   };
 
-  // Build the update package.
-  TemporaryFile zip_file;
-  BuildUpdatePackage(entries, zip_file.release());
+  // "2\nstash " + block3_hash + " 2,2,3"
+  std::string last_command_content = "2\n" + transfer_list_fail[kTransferListHeaderLines + 2];
 
-  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);
+  RunBlockImageUpdate(false, entries, image_file_, "");
 
   // 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 last_command_actual;
+  ASSERT_TRUE(android::base::ReadFileToString(last_command_file_, &last_command_actual));
+  EXPECT_EQ(last_command_content, last_command_actual);
+
   std::string updated_contents;
-  ASSERT_TRUE(android::base::ReadFileToString(update_file.path, &updated_contents));
+  ASSERT_TRUE(android::base::ReadFileToString(image_file_, &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);
+  // "Resume" the update. Expect the first 'move' to be skipped but the second 'move' to be
+  // executed. Note that we intentionally reset the image file.
+  entries["transfer_list"] = android::base::Join(transfer_list_continue, '\n');
+  ASSERT_TRUE(android::base::WriteStringToFile(block1 + block2 + block3, image_file_));
+  RunBlockImageUpdate(false, entries, image_file_, "t");
 
-  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
-  CloseArchive(handle);
+  ASSERT_TRUE(android::base::ReadFileToString(image_file_, &updated_contents));
+  ASSERT_EQ(block1 + block2 + block1, updated_contents);
 }
 
 TEST_F(UpdaterTest, last_command_update_unresumable) {
-  std::string last_command_file = Paths::Get().last_command_file();
-
-  std::string block1 = std::string(4096, '1');
-  std::string block2 = std::string(4096, '2');
+  std::string block1(4096, '1');
+  std::string block2(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::vector<std::string> transfer_list_unresumable{
+    // clang-format off
+    "4",
+    "2",
+    "0",
+    "2",
+    "stash " + block1_hash + " 2,0,1",
+    "move " + block2_hash + " 2,1,2 1 2,0,1",
+    // clang-format on
   };
 
-  std::unordered_map<std::string, std::string> entries = {
+  PackageEntries entries{
     { "new_data", "" },
     { "patch_data", "" },
-    { "transfer_list_unresumable", android::base::Join(transfer_list_unresumable, '\n') },
+    { "transfer_list", android::base::Join(transfer_list_unresumable, '\n') },
   };
 
-  // Build the update package.
-  TemporaryFile zip_file;
-  BuildUpdatePackage(entries, zip_file.release());
+  ASSERT_TRUE(android::base::WriteStringToFile(block1 + block1, image_file_));
 
-  MemMapping map;
-  ASSERT_TRUE(map.MapFile(zip_file.path));
-  ZipArchiveHandle handle;
-  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));
+  std::string last_command_content = "0\n" + transfer_list_unresumable[kTransferListHeaderLines];
+  ASSERT_TRUE(android::base::WriteStringToFile(last_command_content, last_command_file_));
 
-  // 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;
+  RunBlockImageUpdate(false, entries, image_file_, "");
 
-  // 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);
+  // The last_command_file will be deleted if the update encounters an unresumable failure later.
+  ASSERT_EQ(-1, access(last_command_file_.c_str(), R_OK));
 }
 
 TEST_F(UpdaterTest, last_command_verify) {
-  std::string last_command_file = Paths::Get().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(4096, '1');
+  std::string block2(4096, '2');
+  std::string block3(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 = {
+  std::vector<std::string> transfer_list_verify{
+    // clang-format off
     "4",
     "2",
     "0",
@@ -878,55 +798,33 @@
     "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",
+    // clang-format on
   };
 
-  std::unordered_map<std::string, std::string> entries = {
+  PackageEntries entries{
     { "new_data", "" },
     { "patch_data", "" },
-    { "transfer_list_verify", android::base::Join(transfer_list_verify, '\n') },
+    { "transfer_list", android::base::Join(transfer_list_verify, '\n') },
   };
 
-  // Build the update package.
-  TemporaryFile zip_file;
-  BuildUpdatePackage(entries, zip_file.release());
+  ASSERT_TRUE(android::base::WriteStringToFile(block1 + block1 + block3, image_file_));
 
-  MemMapping map;
-  ASSERT_TRUE(map.MapFile(zip_file.path));
-  ZipArchiveHandle handle;
-  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));
+  // Last command: "move " + block1_hash + " 2,1,2 1 2,0,1"
+  std::string last_command_content = "2\n" + transfer_list_verify[kTransferListHeaderLines + 2];
 
-  // 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;
+  // First run: expect the verification to succeed and the last_command_file is intact.
+  ASSERT_TRUE(android::base::WriteStringToFile(last_command_content, last_command_file_));
 
-  std::string src_content = block1 + block1 + block3;
-  TemporaryFile update_file;
-  ASSERT_TRUE(android::base::WriteStringToFile(src_content, update_file.path));
+  RunBlockImageUpdate(true, entries, image_file_, "t");
 
-  ASSERT_TRUE(
-      android::base::WriteStringToFile("2\nstash " + block3_hash + " 2,2,3", last_command_file));
+  std::string last_command_actual;
+  ASSERT_TRUE(android::base::ReadFileToString(last_command_file_, &last_command_actual));
+  EXPECT_EQ(last_command_content, last_command_actual);
 
-  // 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);
+  // Second run with a mismatching block image: expect the verification to succeed but
+  // last_command_file to be deleted; because the target blocks in the last command don't have the
+  // expected contents for the second move command.
+  ASSERT_TRUE(android::base::WriteStringToFile(block1 + block2 + block3, image_file_));
+  RunBlockImageUpdate(true, entries, image_file_, "t");
+  ASSERT_EQ(-1, access(last_command_file_.c_str(), R_OK));
 }
diff --git a/tests/testdata/gzipped_source b/tests/testdata/gzipped_source
new file mode 100644
index 0000000..6d425d0
--- /dev/null
+++ b/tests/testdata/gzipped_source
Binary files differ
diff --git a/tests/testdata/gzipped_target b/tests/testdata/gzipped_target
new file mode 100644
index 0000000..5621262
--- /dev/null
+++ b/tests/testdata/gzipped_target
Binary files differ
diff --git a/ui.h b/ui.h
index 3928426..a74b14f 100644
--- a/ui.h
+++ b/ui.h
@@ -134,6 +134,8 @@
 
   // --- menu display ---
 
+  virtual void SetTitle(const std::vector<std::string>& lines) = 0;
+
   // Displays a menu with the given 'headers' and 'items'. The supplied 'key_handler' callback,
   // which is typically bound to Device::HandleMenuKey(), should return the expected action for the
   // given key code and menu visibility (e.g. to move the cursor or to select an item). Caller sets
diff --git a/updater/blockimg.cpp b/updater/blockimg.cpp
index 236644e..4a70b98 100644
--- a/updater/blockimg.cpp
+++ b/updater/blockimg.cpp
@@ -82,7 +82,7 @@
 
 // Parse the last command index of the last update and save the result to |last_command_index|.
 // Return true if we successfully read the index.
-static bool ParseLastCommandFile(int* last_command_index) {
+static bool ParseLastCommandFile(size_t* last_command_index) {
   const std::string& last_command_file = Paths::Get().last_command_file();
   android::base::unique_fd fd(TEMP_FAILURE_RETRY(open(last_command_file.c_str(), O_RDONLY)));
   if (fd == -1) {
@@ -133,7 +133,7 @@
 }
 
 // Update the last executed command index in the last_command_file.
-static bool UpdateLastCommandIndex(int command_index, const std::string& command_string) {
+static bool UpdateLastCommandIndex(size_t command_index, const std::string& command_string) {
   const std::string& last_command_file = Paths::Get().last_command_file();
   std::string last_command_tmp = last_command_file + ".tmp";
   std::string content = std::to_string(command_index) + "\n" + command_string;
@@ -546,7 +546,6 @@
 struct CommandParameters {
     std::vector<std::string> tokens;
     size_t cpos;
-    int cmdindex;
     const char* cmdname;
     const char* cmdline;
     std::string freestash;
@@ -1666,7 +1665,6 @@
     return StringValue("t");
   }
 
-  size_t start = 2;
   if (lines.size() < 4) {
     ErrorAbort(state, kArgsParsingFailure, "too few lines in the transfer list [%zu]",
                lines.size());
@@ -1691,8 +1689,8 @@
 
   params.createdstash = res;
 
-  // When performing an update, save the index and cmdline of the current command into
-  // the last_command_file.
+  // When performing an update, save the index and cmdline of the current command into the
+  // last_command_file.
   // Upon resuming an update, read the saved index first; then
   //   1. In verification mode, check if the 'move' or 'diff' commands before the saved index has
   //      the expected target blocks already. If not, these commands cannot be skipped and we need
@@ -1701,15 +1699,14 @@
   //   2. In update mode, skip all commands before the saved index. Therefore, we can avoid deleting
   //      stashes with duplicate id unintentionally (b/69858743); and also speed up the update.
   // If an update succeeds or is unresumable, delete the last_command_file.
-  int saved_last_command_index;
+  bool skip_executed_command = true;
+  size_t saved_last_command_index;
   if (!ParseLastCommandFile(&saved_last_command_index)) {
     DeleteLastCommandFile();
-    // We failed to parse the last command, set it explicitly to -1.
-    saved_last_command_index = -1;
+    // We failed to parse the last command. Disallow skipping executed commands.
+    skip_executed_command = false;
   }
 
-  start += 2;
-
   // Build a map of the available commands
   std::unordered_map<std::string, const Command*> cmd_map;
   for (size_t i = 0; i < cmdcount; ++i) {
@@ -1722,18 +1719,15 @@
 
   int rc = -1;
 
+  static constexpr size_t kTransferListHeaderLines = 4;
   // Subsequent lines are all individual transfer commands
-  for (size_t i = start; i < lines.size(); i++) {
+  for (size_t i = kTransferListHeaderLines; i < lines.size(); i++) {
     const std::string& line = lines[i];
     if (line.empty()) continue;
 
+    size_t cmdindex = i - kTransferListHeaderLines;
     params.tokens = android::base::Split(line, " ");
     params.cpos = 0;
-    if (i - start > std::numeric_limits<int>::max()) {
-      params.cmdindex = -1;
-    } else {
-      params.cmdindex = i - start;
-    }
     params.cmdname = params.tokens[params.cpos++].c_str();
     params.cmdline = line.c_str();
     params.target_verified = false;
@@ -1756,9 +1750,9 @@
 
     // Skip all commands before the saved last command index when resuming an update, except for
     // "new" command. Because new commands read in the data sequentially.
-    if (params.canwrite && params.cmdindex != -1 && params.cmdindex <= saved_last_command_index &&
+    if (params.canwrite && skip_executed_command && cmdindex <= saved_last_command_index &&
         cmdname != "new") {
-      LOG(INFO) << "Skipping already executed command: " << params.cmdindex
+      LOG(INFO) << "Skipping already executed command: " << cmdindex
                 << ", last executed command for previous update: " << saved_last_command_index;
       continue;
     }
@@ -1768,17 +1762,16 @@
       goto pbiudone;
     }
 
-    // In verify mode, check if the commands before the saved last_command_index have been
-    // executed correctly. If some target blocks have unexpected contents, delete the last command
-    // file so that we will resume the update from the first command in the transfer list.
-    if (!params.canwrite && saved_last_command_index != -1 && params.cmdindex != -1 &&
-        params.cmdindex <= saved_last_command_index) {
+    // In verify mode, check if the commands before the saved last_command_index have been executed
+    // correctly. If some target blocks have unexpected contents, delete the last command file so
+    // that we will resume the update from the first command in the transfer list.
+    if (!params.canwrite && skip_executed_command && cmdindex <= saved_last_command_index) {
       // TODO(xunchang) check that the cmdline of the saved index is correct.
       if ((cmdname == "move" || cmdname == "bsdiff" || cmdname == "imgdiff") &&
           !params.target_verified) {
         LOG(WARNING) << "Previously executed command " << saved_last_command_index << ": "
                      << params.cmdline << " doesn't produce expected target blocks.";
-        saved_last_command_index = -1;
+        skip_executed_command = false;
         DeleteLastCommandFile();
       }
     }
@@ -1789,7 +1782,7 @@
         goto pbiudone;
       }
 
-      if (!UpdateLastCommandIndex(params.cmdindex, params.cmdline)) {
+      if (!UpdateLastCommandIndex(cmdindex, params.cmdline)) {
         LOG(WARNING) << "Failed to update the last command file.";
       }
 
diff --git a/updater_sample/OWNERS b/updater_sample/OWNERS
new file mode 100644
index 0000000..5c1c370
--- /dev/null
+++ b/updater_sample/OWNERS
@@ -0,0 +1,2 @@
+zhaojiac@google.com
+zhomart@google.com
diff --git a/updater_sample/README.md b/updater_sample/README.md
index 95e57db..3f211dd 100644
--- a/updater_sample/README.md
+++ b/updater_sample/README.md
@@ -44,6 +44,10 @@
 with the offset and length. As `payload.bin` itself is already in compressed
 format, the size penalty is marginal.
 
+if `ab_config.force_switch_slot` set true device will boot to the
+updated partition on next reboot; otherwise button "Switch Slot" will
+become active, and user can manually set updated partition as the active slot.
+
 Config files can be generated using `tools/gen_update_config.py`.
 Running `./tools/gen_update_config.py --help` shows usage of the script.
 
@@ -85,10 +89,10 @@
 - [x] Add stop/reset the update
 - [x] Add demo for passing HTTP headers to `UpdateEngine#applyPayload`
 - [x] [Package compatibility check](https://source.android.com/devices/architecture/vintf/match-rules)
-- [ ] Add tests for `MainActivity`
-- [ ] Change partition demo
+- [x] Deferred switch slot demo
+- [ ] Add demo for passing NETWORK_ID to `UpdateEngine#applyPayload`
 - [ ] Verify system partition checksum for package
-- [ ] Add non-A/B updates demo
+- [?] Add non-A/B updates demo
 
 
 ## Running tests
diff --git a/updater_sample/res/layout/activity_main.xml b/updater_sample/res/layout/activity_main.xml
index 7a12d34..7cde42c 100644
--- a/updater_sample/res/layout/activity_main.xml
+++ b/updater_sample/res/layout/activity_main.xml
@@ -111,19 +111,38 @@
                 android:orientation="horizontal">
 
                 <TextView
-                    android:id="@+id/textView"
+                    android:id="@+id/textView3"
                     android:layout_width="wrap_content"
                     android:layout_height="wrap_content"
-                    android:text="Update status:" />
+                    android:text="Updater state:" />
 
                 <TextView
-                    android:id="@+id/textViewStatus"
+                    android:id="@+id/textViewUpdaterState"
                     android:layout_width="wrap_content"
                     android:layout_height="wrap_content"
                     android:layout_marginLeft="8dp"
                     android:text="@string/unknown" />
             </LinearLayout>
 
+            <LinearLayout
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:layout_marginTop="4dp"
+                android:orientation="horizontal">
+
+                <TextView
+                    android:id="@+id/textView"
+                    android:layout_width="wrap_content"
+                    android:layout_height="wrap_content"
+                    android:text="Engine status:" />
+
+                <TextView
+                    android:id="@+id/textViewEngineStatus"
+                    android:layout_width="wrap_content"
+                    android:layout_height="wrap_content"
+                    android:layout_marginLeft="8dp"
+                    android:text="@string/unknown" />
+            </LinearLayout>
 
             <LinearLayout
                 android:layout_width="match_parent"
@@ -135,10 +154,10 @@
                     android:id="@+id/textView2"
                     android:layout_width="wrap_content"
                     android:layout_height="wrap_content"
-                    android:text="Update completion:" />
+                    android:text="Engine error:" />
 
                 <TextView
-                    android:id="@+id/textViewCompletion"
+                    android:id="@+id/textViewEngineErrorCode"
                     android:layout_width="wrap_content"
                     android:layout_height="wrap_content"
                     android:layout_marginLeft="8dp"
@@ -178,6 +197,23 @@
                     android:text="Reset" />
             </LinearLayout>
 
+            <TextView
+                android:id="@+id/textViewUpdateInfo"
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:layout_marginTop="14dp"
+                android:textColor="#777"
+                android:textSize="10sp"
+                android:textStyle="italic"
+                android:text="@string/finish_update_info" />
+
+            <Button
+                android:id="@+id/buttonSwitchSlot"
+                android:layout_width="wrap_content"
+                android:layout_height="match_parent"
+                android:onClick="onSwitchSlotClick"
+                android:text="@string/switch_slot" />
+
         </LinearLayout>
 
     </ScrollView>
diff --git a/updater_sample/res/raw/sample.json b/updater_sample/res/raw/sample.json
index 46fbfa3..f188c23 100644
--- a/updater_sample/res/raw/sample.json
+++ b/updater_sample/res/raw/sample.json
@@ -20,5 +20,10 @@
             }
         ],
         "authorization": "Basic my-secret-token"
+    },
+    "ab_config": {
+        "__": "A/B (seamless) update configurations",
+        "__force_switch_slot": "if set true device will boot to a new slot, otherwise user manually switches slot on the screen",
+        "force_switch_slot": false
     }
 }
diff --git a/updater_sample/res/values/strings.xml b/updater_sample/res/values/strings.xml
index 2b671ee..db4a5dc 100644
--- a/updater_sample/res/values/strings.xml
+++ b/updater_sample/res/values/strings.xml
@@ -18,4 +18,6 @@
     <string name="action_reload">Reload</string>
     <string name="unknown">Unknown</string>
     <string name="close">CLOSE</string>
+    <string name="switch_slot">Switch Slot</string>
+    <string name="finish_update_info">To finish the update press the button below</string>
 </resources>
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdateConfig.java b/updater_sample/src/com/example/android/systemupdatersample/UpdateConfig.java
index 9bdd8b9..1e0fadc 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdateConfig.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdateConfig.java
@@ -71,7 +71,7 @@
             JSONObject meta = o.getJSONObject("ab_streaming_metadata");
             JSONArray propertyFilesJson = meta.getJSONArray("property_files");
             PackageFile[] propertyFiles =
-                new PackageFile[propertyFilesJson.length()];
+                    new PackageFile[propertyFilesJson.length()];
             for (int i = 0; i < propertyFilesJson.length(); i++) {
                 JSONObject p = propertyFilesJson.getJSONObject(i);
                 propertyFiles[i] = new PackageFile(
@@ -87,6 +87,12 @@
                     propertyFiles,
                     authorization);
         }
+
+        // TODO: parse only for A/B updates when non-A/B is implemented
+        JSONObject ab = o.getJSONObject("ab_config");
+        boolean forceSwitchSlot = ab.getBoolean("force_switch_slot");
+        c.mAbConfig = new AbConfig(forceSwitchSlot);
+
         c.mRawJson = json;
         return c;
     }
@@ -109,6 +115,9 @@
     /** metadata is required only for streaming update */
     private StreamingMetadata mAbStreamingMetadata;
 
+    /** A/B update configurations */
+    private AbConfig mAbConfig;
+
     private String mRawJson;
 
     protected UpdateConfig() {
@@ -119,6 +128,7 @@
         this.mUrl = in.readString();
         this.mAbInstallType = in.readInt();
         this.mAbStreamingMetadata = (StreamingMetadata) in.readSerializable();
+        this.mAbConfig = (AbConfig) in.readSerializable();
         this.mRawJson = in.readString();
     }
 
@@ -148,6 +158,10 @@
         return mAbStreamingMetadata;
     }
 
+    public AbConfig getAbConfig() {
+        return mAbConfig;
+    }
+
     /**
      * @return File object for given url
      */
@@ -172,6 +186,7 @@
         dest.writeString(mUrl);
         dest.writeInt(mAbInstallType);
         dest.writeSerializable(mAbStreamingMetadata);
+        dest.writeSerializable(mAbConfig);
         dest.writeString(mRawJson);
     }
 
@@ -185,9 +200,11 @@
         /** defines beginning of update data in archive */
         private PackageFile[] mPropertyFiles;
 
-        /** SystemUpdaterSample receives the authorization token from the OTA server, in addition
+        /**
+         * SystemUpdaterSample receives the authorization token from the OTA server, in addition
          * to the package URL. It passes on the info to update_engine, so that the latter can
-         * fetch the data from the package server directly with the token. */
+         * fetch the data from the package server directly with the token.
+         */
         private String mAuthorization;
 
         public StreamingMetadata(PackageFile[] propertyFiles, String authorization) {
@@ -239,4 +256,27 @@
         }
     }
 
+    /**
+     * A/B (seamless) update configurations.
+     */
+    public static class AbConfig implements Serializable {
+
+        private static final long serialVersionUID = 31044L;
+
+        /**
+         * if set true device will boot to new slot, otherwise user manually
+         * switches slot on the screen.
+         */
+        private boolean mForceSwitchSlot;
+
+        public AbConfig(boolean forceSwitchSlot) {
+            this.mForceSwitchSlot = forceSwitchSlot;
+        }
+
+        public boolean getForceSwitchSlot() {
+            return mForceSwitchSlot;
+        }
+
+    }
+
 }
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
new file mode 100644
index 0000000..c370a4e
--- /dev/null
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
@@ -0,0 +1,394 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.android.systemupdatersample;
+
+import android.content.Context;
+import android.os.UpdateEngine;
+import android.os.UpdateEngineCallback;
+import android.util.Log;
+
+import com.example.android.systemupdatersample.services.PrepareStreamingService;
+import com.example.android.systemupdatersample.util.PayloadSpecs;
+import com.example.android.systemupdatersample.util.UpdateEngineErrorCodes;
+import com.example.android.systemupdatersample.util.UpdateEngineProperties;
+import com.example.android.systemupdatersample.util.UpdaterStates;
+import com.google.common.util.concurrent.AtomicDouble;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.DoubleConsumer;
+import java.util.function.IntConsumer;
+
+/**
+ * Manages the update flow. It has its own state (in memory), separate from
+ * {@link UpdateEngine}'s state. Asynchronously interacts with the {@link UpdateEngine}.
+ */
+public class UpdateManager {
+
+    private static final String TAG = "UpdateManager";
+
+    /** HTTP Header: User-Agent; it will be sent to the server when streaming the payload. */
+    private static final String HTTP_USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+            + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/60.0.3112.113 Safari/537.36";
+
+    private final UpdateEngine mUpdateEngine;
+    private final PayloadSpecs mPayloadSpecs;
+
+    private AtomicInteger mUpdateEngineStatus =
+            new AtomicInteger(UpdateEngine.UpdateStatusConstants.IDLE);
+    private AtomicInteger mEngineErrorCode = new AtomicInteger(UpdateEngineErrorCodes.UNKNOWN);
+    private AtomicDouble mProgress = new AtomicDouble(0);
+
+    private AtomicInteger mState = new AtomicInteger(UpdaterStates.IDLE);
+
+    private final UpdateManager.UpdateEngineCallbackImpl
+            mUpdateEngineCallback = new UpdateManager.UpdateEngineCallbackImpl();
+
+    private PayloadSpec mLastPayloadSpec;
+    private AtomicBoolean mManualSwitchSlotRequired = new AtomicBoolean(true);
+
+    private IntConsumer mOnStateChangeCallback = null;
+    private IntConsumer mOnEngineStatusUpdateCallback = null;
+    private DoubleConsumer mOnProgressUpdateCallback = null;
+    private IntConsumer mOnEngineCompleteCallback = null;
+
+    private final Object mLock = new Object();
+
+    public UpdateManager(UpdateEngine updateEngine, PayloadSpecs payloadSpecs) {
+        this.mUpdateEngine = updateEngine;
+        this.mPayloadSpecs = payloadSpecs;
+    }
+
+    /**
+     * Binds to {@link UpdateEngine}.
+     */
+    public void bind() {
+        this.mUpdateEngine.bind(mUpdateEngineCallback);
+    }
+
+    /**
+     * Unbinds from {@link UpdateEngine}.
+     */
+    public void unbind() {
+        this.mUpdateEngine.unbind();
+    }
+
+    /**
+     * @return a number from {@code 0.0} to {@code 1.0}.
+     */
+    public float getProgress() {
+        return (float) this.mProgress.get();
+    }
+
+    /**
+     * Returns true if manual switching slot is required. Value depends on
+     * the update config {@code ab_config.force_switch_slot}.
+     */
+    public boolean isManualSwitchSlotRequired() {
+        return mManualSwitchSlotRequired.get();
+    }
+
+    /**
+     * Sets SystemUpdaterSample app state change callback. Value of {@code state} will be one
+     * of the values from {@link UpdaterStates}.
+     *
+     * @param onStateChangeCallback a callback with parameter {@code state}.
+     */
+    public void setOnStateChangeCallback(IntConsumer onStateChangeCallback) {
+        synchronized (mLock) {
+            this.mOnStateChangeCallback = onStateChangeCallback;
+        }
+    }
+
+    private Optional<IntConsumer> getOnStateChangeCallback() {
+        synchronized (mLock) {
+            return mOnStateChangeCallback == null
+                    ? Optional.empty()
+                    : Optional.of(mOnStateChangeCallback);
+        }
+    }
+
+    /**
+     * Sets update engine status update callback. Value of {@code status} will
+     * be one of the values from {@link UpdateEngine.UpdateStatusConstants}.
+     *
+     * @param onStatusUpdateCallback a callback with parameter {@code status}.
+     */
+    public void setOnEngineStatusUpdateCallback(IntConsumer onStatusUpdateCallback) {
+        synchronized (mLock) {
+            this.mOnEngineStatusUpdateCallback = onStatusUpdateCallback;
+        }
+    }
+
+    private Optional<IntConsumer> getOnEngineStatusUpdateCallback() {
+        synchronized (mLock) {
+            return mOnEngineStatusUpdateCallback == null
+                    ? Optional.empty()
+                    : Optional.of(mOnEngineStatusUpdateCallback);
+        }
+    }
+
+    /**
+     * Sets update engine payload application complete callback. Value of {@code errorCode} will
+     * be one of the values from {@link UpdateEngine.ErrorCodeConstants}.
+     *
+     * @param onComplete a callback with parameter {@code errorCode}.
+     */
+    public void setOnEngineCompleteCallback(IntConsumer onComplete) {
+        synchronized (mLock) {
+            this.mOnEngineCompleteCallback = onComplete;
+        }
+    }
+
+    private Optional<IntConsumer> getOnEngineCompleteCallback() {
+        synchronized (mLock) {
+            return mOnEngineCompleteCallback == null
+                    ? Optional.empty()
+                    : Optional.of(mOnEngineCompleteCallback);
+        }
+    }
+
+    /**
+     * Sets progress update callback. Progress is a number from {@code 0.0} to {@code 1.0}.
+     *
+     * @param onProgressCallback a callback with parameter {@code progress}.
+     */
+    public void setOnProgressUpdateCallback(DoubleConsumer onProgressCallback) {
+        synchronized (mLock) {
+            this.mOnProgressUpdateCallback = onProgressCallback;
+        }
+    }
+
+    private Optional<DoubleConsumer> getOnProgressUpdateCallback() {
+        synchronized (mLock) {
+            return mOnProgressUpdateCallback == null
+                    ? Optional.empty()
+                    : Optional.of(mOnProgressUpdateCallback);
+        }
+    }
+
+    /**
+     * Updates {@link this.mState} and if state is changed,
+     * it also notifies {@link this.mOnStateChangeCallback}.
+     */
+    private void setUpdaterState(int updaterState) {
+        int previousState = mState.get();
+        mState.set(updaterState);
+        if (previousState != updaterState) {
+            getOnStateChangeCallback().ifPresent(callback -> callback.accept(updaterState));
+        }
+    }
+
+    /**
+     * Requests update engine to stop any ongoing update. If an update has been applied,
+     * leave it as is.
+     *
+     * <p>Sometimes it's possible that the
+     * update engine would throw an error when the method is called, and the only way to
+     * handle it is to catch the exception.</p>
+     */
+    public void cancelRunningUpdate() {
+        try {
+            mUpdateEngine.cancel();
+            setUpdaterState(UpdaterStates.IDLE);
+        } catch (Exception e) {
+            Log.w(TAG, "UpdateEngine failed to stop the ongoing update", e);
+        }
+    }
+
+    /**
+     * Resets update engine to IDLE state. If an update has been applied it reverts it.
+     *
+     * <p>Sometimes it's possible that the
+     * update engine would throw an error when the method is called, and the only way to
+     * handle it is to catch the exception.</p>
+     */
+    public void resetUpdate() {
+        try {
+            mUpdateEngine.resetStatus();
+            setUpdaterState(UpdaterStates.IDLE);
+        } catch (Exception e) {
+            Log.w(TAG, "UpdateEngine failed to reset the update", e);
+        }
+    }
+
+    /**
+     * Applies the given update.
+     *
+     * <p>UpdateEngine works asynchronously. This method doesn't wait until
+     * end of the update.</p>
+     */
+    public void applyUpdate(Context context, UpdateConfig config) {
+        mEngineErrorCode.set(UpdateEngineErrorCodes.UNKNOWN);
+        setUpdaterState(UpdaterStates.RUNNING);
+
+        if (!config.getAbConfig().getForceSwitchSlot()) {
+            mManualSwitchSlotRequired.set(true);
+        } else {
+            mManualSwitchSlotRequired.set(false);
+        }
+
+        if (config.getInstallType() == UpdateConfig.AB_INSTALL_TYPE_NON_STREAMING) {
+            applyAbNonStreamingUpdate(config);
+        } else {
+            applyAbStreamingUpdate(context, config);
+        }
+    }
+
+    private void applyAbNonStreamingUpdate(UpdateConfig config) {
+        List<String> extraProperties = prepareExtraProperties(config);
+
+        PayloadSpec payload;
+        try {
+            payload = mPayloadSpecs.forNonStreaming(config.getUpdatePackageFile());
+        } catch (IOException e) {
+            Log.e(TAG, "Error creating payload spec", e);
+            setUpdaterState(UpdaterStates.ERROR);
+            return;
+        }
+        updateEngineApplyPayload(payload, extraProperties);
+    }
+
+    private void applyAbStreamingUpdate(Context context, UpdateConfig config) {
+        List<String> extraProperties = prepareExtraProperties(config);
+
+        Log.d(TAG, "Starting PrepareStreamingService");
+        PrepareStreamingService.startService(context, config, (code, payloadSpec) -> {
+            if (code == PrepareStreamingService.RESULT_CODE_SUCCESS) {
+                extraProperties.add("USER_AGENT=" + HTTP_USER_AGENT);
+                config.getStreamingMetadata()
+                        .getAuthorization()
+                        .ifPresent(s -> extraProperties.add("AUTHORIZATION=" + s));
+                updateEngineApplyPayload(payloadSpec, extraProperties);
+            } else {
+                Log.e(TAG, "PrepareStreamingService failed, result code is " + code);
+                setUpdaterState(UpdaterStates.ERROR);
+            }
+        });
+    }
+
+    private List<String> prepareExtraProperties(UpdateConfig config) {
+        List<String> extraProperties = new ArrayList<>();
+
+        if (!config.getAbConfig().getForceSwitchSlot()) {
+            // Disable switch slot on reboot, which is enabled by default.
+            // User will enable it manually by clicking "Switch Slot" button on the screen.
+            extraProperties.add(UpdateEngineProperties.PROPERTY_DISABLE_SWITCH_SLOT_ON_REBOOT);
+        }
+        return extraProperties;
+    }
+
+    /**
+     * Applies given payload.
+     *
+     * <p>UpdateEngine works asynchronously. This method doesn't wait until
+     * end of the update.</p>
+     *
+     * <p>It's possible that the update engine throws a generic error, such as upon seeing invalid
+     * payload properties (which come from OTA packages), or failing to set up the network
+     * with the given id.</p>
+     *
+     * @param payloadSpec contains url, offset and size to {@code PAYLOAD_BINARY_FILE_NAME}
+     * @param extraProperties additional properties to pass to {@link UpdateEngine#applyPayload}
+     */
+    private void updateEngineApplyPayload(PayloadSpec payloadSpec, List<String> extraProperties) {
+        mLastPayloadSpec = payloadSpec;
+
+        ArrayList<String> properties = new ArrayList<>(payloadSpec.getProperties());
+        if (extraProperties != null) {
+            properties.addAll(extraProperties);
+        }
+        try {
+            mUpdateEngine.applyPayload(
+                    payloadSpec.getUrl(),
+                    payloadSpec.getOffset(),
+                    payloadSpec.getSize(),
+                    properties.toArray(new String[0]));
+        } catch (Exception e) {
+            Log.e(TAG, "UpdateEngine failed to apply the update", e);
+            setUpdaterState(UpdaterStates.ERROR);
+        }
+    }
+
+    /**
+     * Sets the new slot that has the updated partitions as the active slot,
+     * which device will boot into next time.
+     * This method is only supposed to be called after the payload is applied.
+     *
+     * Invoking {@link UpdateEngine#applyPayload} with the same payload url, offset, size
+     * and payload metadata headers doesn't trigger new update. It can be used to just switch
+     * active A/B slot.
+     *
+     * {@link UpdateEngine#applyPayload} might take several seconds to finish, and it will
+     * invoke callbacks {@link this#onStatusUpdate} and {@link this#onPayloadApplicationComplete)}.
+     */
+    public void setSwitchSlotOnReboot() {
+        Log.d(TAG, "setSwitchSlotOnReboot invoked");
+        List<String> extraProperties = new ArrayList<>();
+        // PROPERTY_SKIP_POST_INSTALL should be passed on to skip post-installation hooks.
+        extraProperties.add(UpdateEngineProperties.PROPERTY_SKIP_POST_INSTALL);
+        // It sets property SWITCH_SLOT_ON_REBOOT=1 by default.
+        // HTTP headers are not required, UpdateEngine is not expected to stream payload.
+        updateEngineApplyPayload(mLastPayloadSpec, extraProperties);
+    }
+
+    private void onStatusUpdate(int status, float progress) {
+        int previousStatus = mUpdateEngineStatus.get();
+        mUpdateEngineStatus.set(status);
+        mProgress.set(progress);
+
+        getOnProgressUpdateCallback().ifPresent(callback -> callback.accept(progress));
+
+        if (previousStatus != status) {
+            getOnEngineStatusUpdateCallback().ifPresent(callback -> callback.accept(status));
+        }
+    }
+
+    private void onPayloadApplicationComplete(int errorCode) {
+        Log.d(TAG, "onPayloadApplicationComplete invoked, errorCode=" + errorCode);
+        mEngineErrorCode.set(errorCode);
+        if (errorCode == UpdateEngine.ErrorCodeConstants.SUCCESS
+                || errorCode == UpdateEngineErrorCodes.UPDATED_BUT_NOT_ACTIVE) {
+            setUpdaterState(UpdaterStates.FINISHED);
+        } else if (errorCode != UpdateEngineErrorCodes.USER_CANCELLED) {
+            setUpdaterState(UpdaterStates.ERROR);
+        }
+
+        getOnEngineCompleteCallback()
+                .ifPresent(callback -> callback.accept(errorCode));
+    }
+
+    /**
+     * Helper class to delegate {@code update_engine} callbacks to UpdateManager
+     */
+    class UpdateEngineCallbackImpl extends UpdateEngineCallback {
+        @Override
+        public void onStatusUpdate(int status, float percent) {
+            UpdateManager.this.onStatusUpdate(status, percent);
+        }
+
+        @Override
+        public void onPayloadApplicationComplete(int errorCode) {
+            UpdateManager.this.onPayloadApplicationComplete(errorCode);
+        }
+    }
+
+}
diff --git a/updater_sample/src/com/example/android/systemupdatersample/services/PrepareStreamingService.java b/updater_sample/src/com/example/android/systemupdatersample/services/PrepareStreamingService.java
index 222bb0a..ac6e223 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/services/PrepareStreamingService.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/services/PrepareStreamingService.java
@@ -116,6 +116,8 @@
                 PackageFiles.PAYLOAD_PROPERTIES_FILE_NAME
             );
 
+    private final PayloadSpecs mPayloadSpecs = new PayloadSpecs();
+
     @Override
     protected void onHandleIntent(Intent intent) {
         Log.d(TAG, "On handle intent is called");
@@ -137,7 +139,7 @@
      * 3. Checks OTA package compatibility with the device.
      * 4. Constructs {@link PayloadSpec} for streaming update.
      */
-    private static PayloadSpec execute(UpdateConfig config)
+    private PayloadSpec execute(UpdateConfig config)
             throws IOException, PreparationFailedException {
 
         downloadPreStreamingFiles(config, OTA_PACKAGE_DIR);
@@ -164,7 +166,7 @@
             }
         }
 
-        return PayloadSpecs.forStreaming(config.getUrl(),
+        return mPayloadSpecs.forStreaming(config.getUrl(),
                 payloadBinary.get().getOffset(),
                 payloadBinary.get().getSize(),
                 Paths.get(OTA_PACKAGE_DIR, PAYLOAD_PROPERTIES_FILE_NAME).toFile());
@@ -176,7 +178,7 @@
      * in directory {@code dir}.
      * @throws IOException when can't download a file
      */
-    private static void downloadPreStreamingFiles(UpdateConfig config, String dir)
+    private void downloadPreStreamingFiles(UpdateConfig config, String dir)
             throws IOException {
         Log.d(TAG, "Deleting existing files from " + dir);
         for (String file : PRE_STREAMING_FILES_SET) {
@@ -200,7 +202,7 @@
      * @param file physical location of {@link PackageFiles#COMPATIBILITY_ZIP_FILE_NAME}
      * @return true if OTA package is compatible with this device
      */
-    private static boolean verifyPackageCompatibility(File file) {
+    private boolean verifyPackageCompatibility(File file) {
         try {
             return RecoverySystem.verifyPackageCompatibility(file);
         } catch (IOException e) {
diff --git a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
index 1708256..9983fe3 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
@@ -18,10 +18,10 @@
 
 import android.app.Activity;
 import android.app.AlertDialog;
+import android.graphics.Color;
 import android.os.Build;
 import android.os.Bundle;
 import android.os.UpdateEngine;
-import android.os.UpdateEngineCallback;
 import android.util.Log;
 import android.view.View;
 import android.widget.ArrayAdapter;
@@ -29,21 +29,17 @@
 import android.widget.ProgressBar;
 import android.widget.Spinner;
 import android.widget.TextView;
-import android.widget.Toast;
 
-import com.example.android.systemupdatersample.PayloadSpec;
 import com.example.android.systemupdatersample.R;
 import com.example.android.systemupdatersample.UpdateConfig;
-import com.example.android.systemupdatersample.services.PrepareStreamingService;
+import com.example.android.systemupdatersample.UpdateManager;
 import com.example.android.systemupdatersample.util.PayloadSpecs;
 import com.example.android.systemupdatersample.util.UpdateConfigs;
 import com.example.android.systemupdatersample.util.UpdateEngineErrorCodes;
 import com.example.android.systemupdatersample.util.UpdateEngineStatuses;
+import com.example.android.systemupdatersample.util.UpdaterStates;
 
-import java.io.IOException;
-import java.util.ArrayList;
 import java.util.List;
-import java.util.concurrent.atomic.AtomicInteger;
 
 /**
  * UI for SystemUpdaterSample app.
@@ -52,10 +48,6 @@
 
     private static final String TAG = "MainActivity";
 
-    /** HTTP Header: User-Agent; it will be sent to the server when streaming the payload. */
-    private static final String HTTP_USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
-            + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/60.0.3112.113 Safari/537.36";
-
     private TextView mTextViewBuild;
     private Spinner mSpinnerConfigs;
     private TextView mTextViewConfigsDirHint;
@@ -64,19 +56,16 @@
     private Button mButtonStop;
     private Button mButtonReset;
     private ProgressBar mProgressBar;
-    private TextView mTextViewStatus;
-    private TextView mTextViewCompletion;
+    private TextView mTextViewUpdaterState;
+    private TextView mTextViewEngineStatus;
+    private TextView mTextViewEngineErrorCode;
+    private TextView mTextViewUpdateInfo;
+    private Button mButtonSwitchSlot;
 
     private List<UpdateConfig> mConfigs;
-    private AtomicInteger mUpdateEngineStatus =
-            new AtomicInteger(UpdateEngine.UpdateStatusConstants.IDLE);
 
-    /**
-     * Listen to {@code update_engine} events.
-     */
-    private UpdateEngineCallbackImpl mUpdateEngineCallback = new UpdateEngineCallbackImpl();
-
-    private final UpdateEngine mUpdateEngine = new UpdateEngine();
+    private final UpdateManager mUpdateManager =
+            new UpdateManager(new UpdateEngine(), new PayloadSpecs());
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
@@ -91,23 +80,43 @@
         this.mButtonStop = findViewById(R.id.buttonStop);
         this.mButtonReset = findViewById(R.id.buttonReset);
         this.mProgressBar = findViewById(R.id.progressBar);
-        this.mTextViewStatus = findViewById(R.id.textViewStatus);
-        this.mTextViewCompletion = findViewById(R.id.textViewCompletion);
+        this.mTextViewUpdaterState = findViewById(R.id.textViewUpdaterState);
+        this.mTextViewEngineStatus = findViewById(R.id.textViewEngineStatus);
+        this.mTextViewEngineErrorCode = findViewById(R.id.textViewEngineErrorCode);
+        this.mTextViewUpdateInfo = findViewById(R.id.textViewUpdateInfo);
+        this.mButtonSwitchSlot = findViewById(R.id.buttonSwitchSlot);
 
         this.mTextViewConfigsDirHint.setText(UpdateConfigs.getConfigsRoot(this));
 
         uiReset();
         loadUpdateConfigs();
 
-        this.mUpdateEngine.bind(mUpdateEngineCallback);
+        this.mUpdateManager.setOnStateChangeCallback(this::onUpdaterStateChange);
+        this.mUpdateManager.setOnEngineStatusUpdateCallback(this::onEngineStatusUpdate);
+        this.mUpdateManager.setOnEngineCompleteCallback(this::onEnginePayloadApplicationComplete);
+        this.mUpdateManager.setOnProgressUpdateCallback(this::onProgressUpdate);
     }
 
     @Override
     protected void onDestroy() {
-        this.mUpdateEngine.unbind();
+        this.mUpdateManager.setOnEngineStatusUpdateCallback(null);
+        this.mUpdateManager.setOnProgressUpdateCallback(null);
+        this.mUpdateManager.setOnEngineCompleteCallback(null);
         super.onDestroy();
     }
 
+    @Override
+    protected void onResume() {
+        super.onResume();
+        this.mUpdateManager.bind();
+    }
+
+    @Override
+    protected void onPause() {
+        this.mUpdateManager.unbind();
+        super.onPause();
+    }
+
     /**
      * reload button is clicked
      */
@@ -137,7 +146,8 @@
                 .setIcon(android.R.drawable.ic_dialog_alert)
                 .setPositiveButton(android.R.string.ok, (dialog, whichButton) -> {
                     uiSetUpdating();
-                    applyUpdate(getSelectedConfig());
+                    uiResetEngineText();
+                    mUpdateManager.applyUpdate(this, getSelectedConfig());
                 })
                 .setNegativeButton(android.R.string.cancel, null)
                 .show();
@@ -152,7 +162,7 @@
                 .setMessage("Do you really want to cancel running update?")
                 .setIcon(android.R.drawable.ic_dialog_alert)
                 .setPositiveButton(android.R.string.ok, (dialog, whichButton) -> {
-                    stopRunningUpdate();
+                    mUpdateManager.cancelRunningUpdate();
                 })
                 .setNegativeButton(android.R.string.cancel, null).show();
     }
@@ -167,36 +177,48 @@
                         + " and restore old version?")
                 .setIcon(android.R.drawable.ic_dialog_alert)
                 .setPositiveButton(android.R.string.ok, (dialog, whichButton) -> {
-                    resetUpdate();
+                    mUpdateManager.resetUpdate();
                 })
                 .setNegativeButton(android.R.string.cancel, null).show();
     }
 
     /**
-     * Invoked when anything changes. The value of {@code status} will
-     * be one of the values from {@link UpdateEngine.UpdateStatusConstants},
-     * and {@code percent} will be from {@code 0.0} to {@code 1.0}.
+     * switch slot button clicked
      */
-    private void onStatusUpdate(int status, float percent) {
-        mProgressBar.setProgress((int) (100 * percent));
-        if (mUpdateEngineStatus.get() != status) {
-            mUpdateEngineStatus.set(status);
-            runOnUiThread(() -> {
-                Log.e("UpdateEngine", "StatusUpdate - status="
-                        + UpdateEngineStatuses.getStatusText(status)
-                        + "/" + status);
-                setUiStatus(status);
-                Toast.makeText(this, "Update Status changed", Toast.LENGTH_LONG)
-                        .show();
-                if (status != UpdateEngine.UpdateStatusConstants.IDLE) {
-                    Log.d(TAG, "status changed, setting ui to updating mode");
-                    uiSetUpdating();
-                } else {
-                    Log.d(TAG, "status changed, resetting ui");
-                    uiReset();
-                }
-            });
-        }
+    public void onSwitchSlotClick(View view) {
+        mUpdateManager.setSwitchSlotOnReboot();
+    }
+
+    /**
+     * Invoked when SystemUpdaterSample app state changes.
+     * Value of {@code state} will be one of the
+     * values from {@link UpdaterStates}.
+     */
+    private void onUpdaterStateChange(int state) {
+        Log.i(TAG, "onUpdaterStateChange invoked state=" + state);
+        runOnUiThread(() -> {
+            setUiUpdaterState(state);
+        });
+    }
+
+    /**
+     * Invoked when {@link UpdateEngine} status changes. Value of {@code status} will
+     * be one of the values from {@link UpdateEngine.UpdateStatusConstants}.
+     */
+    private void onEngineStatusUpdate(int status) {
+        runOnUiThread(() -> {
+            Log.e(TAG, "StatusUpdate - status="
+                    + UpdateEngineStatuses.getStatusText(status)
+                    + "/" + status);
+            if (status == UpdateEngine.UpdateStatusConstants.IDLE) {
+                Log.d(TAG, "status changed, resetting ui");
+                uiReset();
+            } else {
+                Log.d(TAG, "status changed, setting ui to updating mode");
+                uiSetUpdating();
+            }
+            setUiEngineStatus(status);
+        });
     }
 
     /**
@@ -204,20 +226,33 @@
      * unsuccessfully. The value of {@code errorCode} will be one of the
      * values from {@link UpdateEngine.ErrorCodeConstants}.
      */
-    private void onPayloadApplicationComplete(int errorCode) {
-        final String state = UpdateEngineErrorCodes.isUpdateSucceeded(errorCode)
+    private void onEnginePayloadApplicationComplete(int errorCode) {
+        final String completionState = UpdateEngineErrorCodes.isUpdateSucceeded(errorCode)
                 ? "SUCCESS"
                 : "FAILURE";
         runOnUiThread(() -> {
-            Log.i("UpdateEngine",
+            Log.i(TAG,
                     "Completed - errorCode="
                     + UpdateEngineErrorCodes.getCodeName(errorCode) + "/" + errorCode
-                    + " " + state);
-            Toast.makeText(this, "Update completed", Toast.LENGTH_LONG).show();
-            setUiCompletion(errorCode);
+                    + " " + completionState);
+            setUiEngineErrorCode(errorCode);
+            if (errorCode == UpdateEngineErrorCodes.UPDATED_BUT_NOT_ACTIVE) {
+                // if update was successfully applied.
+                if (mUpdateManager.isManualSwitchSlotRequired()) {
+                    // Show "Switch Slot" button.
+                    uiShowSwitchSlotInfo();
+                }
+            }
         });
     }
 
+    /**
+     * Invoked when update progress changes.
+     */
+    private void onProgressUpdate(double progress) {
+        mProgressBar.setProgress((int) (100 * progress));
+    }
+
     /** resets ui */
     private void uiReset() {
         mTextViewBuild.setText(Build.DISPLAY);
@@ -229,8 +264,13 @@
         mProgressBar.setProgress(0);
         mProgressBar.setEnabled(false);
         mProgressBar.setVisibility(ProgressBar.INVISIBLE);
-        mTextViewStatus.setText(R.string.unknown);
-        mTextViewCompletion.setText(R.string.unknown);
+        uiHideSwitchSlotInfo();
+    }
+
+    private void uiResetEngineText() {
+        mTextViewEngineStatus.setText(R.string.unknown);
+        mTextViewEngineErrorCode.setText(R.string.unknown);
+        // Note: Do not reset mTextViewUpdaterState; UpdateManager notifies properly.
     }
 
     /** sets ui updating mode */
@@ -245,6 +285,16 @@
         mProgressBar.setVisibility(ProgressBar.VISIBLE);
     }
 
+    private void uiShowSwitchSlotInfo() {
+        mButtonSwitchSlot.setEnabled(true);
+        mTextViewUpdateInfo.setTextColor(Color.parseColor("#777777"));
+    }
+
+    private void uiHideSwitchSlotInfo() {
+        mTextViewUpdateInfo.setTextColor(Color.parseColor("#AAAAAA"));
+        mButtonSwitchSlot.setEnabled(false);
+    }
+
     /**
      * loads json configurations from configs dir that is defined in {@link UpdateConfigs}.
      */
@@ -256,20 +306,25 @@
     /**
      * @param status update engine status code
      */
-    private void setUiStatus(int status) {
+    private void setUiEngineStatus(int status) {
         String statusText = UpdateEngineStatuses.getStatusText(status);
-        mTextViewStatus.setText(statusText + "/" + status);
+        mTextViewEngineStatus.setText(statusText + "/" + status);
     }
 
     /**
      * @param errorCode update engine error code
      */
-    private void setUiCompletion(int errorCode) {
-        final String state = UpdateEngineErrorCodes.isUpdateSucceeded(errorCode)
-                ? "SUCCESS"
-                : "FAILURE";
+    private void setUiEngineErrorCode(int errorCode) {
         String errorText = UpdateEngineErrorCodes.getCodeName(errorCode);
-        mTextViewCompletion.setText(state + " " + errorText + "/" + errorCode);
+        mTextViewEngineErrorCode.setText(errorText + "/" + errorCode);
+    }
+
+    /**
+     * @param state updater sample state
+     */
+    private void setUiUpdaterState(int state) {
+        String stateText = UpdaterStates.getStateText(state);
+        mTextViewUpdaterState.setText(stateText + "/" + state);
     }
 
     private void loadConfigsToSpinner(List<UpdateConfig> configs) {
@@ -286,108 +341,4 @@
         return mConfigs.get(mSpinnerConfigs.getSelectedItemPosition());
     }
 
-    /**
-     * Applies the given update
-     */
-    private void applyUpdate(final UpdateConfig config) {
-        if (config.getInstallType() == UpdateConfig.AB_INSTALL_TYPE_NON_STREAMING) {
-            PayloadSpec payload;
-            try {
-                payload = PayloadSpecs.forNonStreaming(config.getUpdatePackageFile());
-            } catch (IOException e) {
-                Log.e(TAG, "Error creating payload spec", e);
-                Toast.makeText(this, "Error creating payload spec", Toast.LENGTH_LONG)
-                        .show();
-                return;
-            }
-            updateEngineApplyPayload(payload, null);
-        } else {
-            Log.d(TAG, "Starting PrepareStreamingService");
-            PrepareStreamingService.startService(this, config, (code, payloadSpec) -> {
-                if (code == PrepareStreamingService.RESULT_CODE_SUCCESS) {
-                    List<String> extraProperties = new ArrayList<>();
-                    extraProperties.add("USER_AGENT=" + HTTP_USER_AGENT);
-                    config.getStreamingMetadata()
-                            .getAuthorization()
-                            .ifPresent(s -> extraProperties.add("AUTHORIZATION=" + s));
-                    updateEngineApplyPayload(payloadSpec, extraProperties);
-                } else {
-                    Log.e(TAG, "PrepareStreamingService failed, result code is " + code);
-                    Toast.makeText(
-                            MainActivity.this,
-                            "PrepareStreamingService failed, result code is " + code,
-                            Toast.LENGTH_LONG).show();
-                }
-            });
-        }
-    }
-
-    /**
-     * Applies given payload.
-     *
-     * UpdateEngine works asynchronously. This method doesn't wait until
-     * end of the update.
-     *
-     * @param payloadSpec contains url, offset and size to {@code PAYLOAD_BINARY_FILE_NAME}
-     * @param extraProperties additional properties to pass to {@link UpdateEngine#applyPayload}
-     */
-    private void updateEngineApplyPayload(PayloadSpec payloadSpec, List<String> extraProperties) {
-        ArrayList<String> properties = new ArrayList<>(payloadSpec.getProperties());
-        if (extraProperties != null) {
-            properties.addAll(extraProperties);
-        }
-        try {
-            mUpdateEngine.applyPayload(
-                    payloadSpec.getUrl(),
-                    payloadSpec.getOffset(),
-                    payloadSpec.getSize(),
-                    properties.toArray(new String[0]));
-        } catch (Exception e) {
-            Log.e(TAG, "UpdateEngine failed to apply the update", e);
-            Toast.makeText(
-                    this,
-                    "UpdateEngine failed to apply the update",
-                    Toast.LENGTH_LONG).show();
-        }
-    }
-
-    /**
-     * Requests update engine to stop any ongoing update. If an update has been applied,
-     * leave it as is.
-     */
-    private void stopRunningUpdate() {
-        try {
-            mUpdateEngine.cancel();
-        } catch (Exception e) {
-            Log.w(TAG, "UpdateEngine failed to stop the ongoing update", e);
-        }
-    }
-
-    /**
-     * Resets update engine to IDLE state. Requests to cancel any onging update, or to revert if an
-     * update has been applied.
-     */
-    private void resetUpdate() {
-        try {
-            mUpdateEngine.resetStatus();
-        } catch (Exception e) {
-            Log.w(TAG, "UpdateEngine failed to reset the update", e);
-        }
-    }
-
-    /**
-     * Helper class to delegate {@code update_engine} callbacks to MainActivity
-     */
-    class UpdateEngineCallbackImpl extends UpdateEngineCallback {
-        @Override
-        public void onStatusUpdate(int status, float percent) {
-            MainActivity.this.onStatusUpdate(status, percent);
-        }
-
-        @Override
-        public void onPayloadApplicationComplete(int errorCode) {
-            MainActivity.this.onPayloadApplicationComplete(errorCode);
-        }
-    }
-
 }
diff --git a/updater_sample/src/com/example/android/systemupdatersample/util/PayloadSpecs.java b/updater_sample/src/com/example/android/systemupdatersample/util/PayloadSpecs.java
index 4db448a..f062317 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/util/PayloadSpecs.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/util/PayloadSpecs.java
@@ -32,7 +32,9 @@
 import java.util.zip.ZipFile;
 
 /** The helper class that creates {@link PayloadSpec}. */
-public final class PayloadSpecs {
+public class PayloadSpecs {
+
+    public PayloadSpecs() {}
 
     /**
      * The payload PAYLOAD_ENTRY is stored in the zip package to comply with the Android OTA package
@@ -43,7 +45,7 @@
      * zip file. So we enumerate the entries to identify the offset of the payload file.
      * http://developer.android.com/reference/java/util/zip/ZipFile.html#entries()
      */
-    public static PayloadSpec forNonStreaming(File packageFile) throws IOException {
+    public PayloadSpec forNonStreaming(File packageFile) throws IOException {
         boolean payloadFound = false;
         long payloadOffset = 0;
         long payloadSize = 0;
@@ -100,7 +102,7 @@
     /**
      * Creates a {@link PayloadSpec} for streaming update.
      */
-    public static PayloadSpec forStreaming(String updateUrl,
+    public PayloadSpec forStreaming(String updateUrl,
                                            long offset,
                                            long size,
                                            File propertiesFile) throws IOException {
@@ -115,7 +117,7 @@
     /**
      * Converts an {@link PayloadSpec} to a string.
      */
-    public static String toString(PayloadSpec payloadSpec) {
+    public String specToString(PayloadSpec payloadSpec) {
         return "<PayloadSpec url=" + payloadSpec.getUrl()
                 + ", offset=" + payloadSpec.getOffset()
                 + ", size=" + payloadSpec.getSize()
@@ -124,6 +126,4 @@
                 + ">";
     }
 
-    private PayloadSpecs() {}
-
 }
diff --git a/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineErrorCodes.java b/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineErrorCodes.java
index 6d319c5..7d55ff8 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineErrorCodes.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineErrorCodes.java
@@ -34,7 +34,9 @@
     * Error code from the update engine. Values must agree with the ones in
     * system/update_engine/common/error_code.h.
     */
+    public static final int UNKNOWN = -1;
     public static final int UPDATED_BUT_NOT_ACTIVE = 52;
+    public static final int USER_CANCELLED = 48;
 
     private static final SparseArray<String> CODE_TO_NAME_MAP = new SparseArray<>();
 
@@ -60,7 +62,7 @@
      * Completion codes returned by update engine indicating that the update
      * was successfully applied.
      */
-    private static final Set<Integer> SUCCEEDED_COMPLETION_CODES = new HashSet<Integer>(
+    private static final Set<Integer> SUCCEEDED_COMPLETION_CODES = new HashSet<>(
             Arrays.asList(UpdateEngine.ErrorCodeConstants.SUCCESS,
                     // UPDATED_BUT_NOT_ACTIVE is returned when the payload is
                     // successfully applied but the
diff --git a/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineProperties.java b/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineProperties.java
new file mode 100644
index 0000000..e368f14
--- /dev/null
+++ b/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineProperties.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.android.systemupdatersample.util;
+
+/**
+ * Utility class for properties that will be passed to {@code UpdateEngine#applyPayload}.
+ */
+public final class UpdateEngineProperties {
+
+    /**
+     * The property indicating that the update engine should not switch slot
+     * when the device reboots.
+     */
+    public static final String PROPERTY_DISABLE_SWITCH_SLOT_ON_REBOOT = "SWITCH_SLOT_ON_REBOOT=0";
+
+    /**
+     * The property to skip post-installation.
+     * https://source.android.com/devices/tech/ota/ab/#post-installation
+     */
+    public static final String PROPERTY_SKIP_POST_INSTALL = "RUN_POST_INSTALL=0";
+
+    private UpdateEngineProperties() {}
+}
diff --git a/updater_sample/src/com/example/android/systemupdatersample/util/UpdaterStates.java b/updater_sample/src/com/example/android/systemupdatersample/util/UpdaterStates.java
new file mode 100644
index 0000000..fc20a79
--- /dev/null
+++ b/updater_sample/src/com/example/android/systemupdatersample/util/UpdaterStates.java
@@ -0,0 +1,50 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.android.systemupdatersample.util;
+
+import android.util.SparseArray;
+
+/**
+ * SystemUpdaterSample app state.
+ */
+public class UpdaterStates {
+
+    public static final int IDLE = 0;
+    public static final int ERROR = 1;
+    public static final int RUNNING = 2;
+    public static final int PAUSED = 3;
+    public static final int FINISHED = 4;
+
+    private static final SparseArray<String> STATE_MAP = new SparseArray<>();
+
+    static {
+        STATE_MAP.put(0, "IDLE");
+        STATE_MAP.put(1, "ERROR");
+        STATE_MAP.put(2, "RUNNING");
+        STATE_MAP.put(3, "PAUSED");
+        STATE_MAP.put(4, "FINISHED");
+    }
+
+    /**
+     * converts status code to status name
+     */
+    public static String getStateText(int state) {
+        return STATE_MAP.get(state);
+    }
+
+    private UpdaterStates() {}
+}
diff --git a/updater_sample/tests/Android.mk b/updater_sample/tests/Android.mk
index a1a4664..9aec372 100644
--- a/updater_sample/tests/Android.mk
+++ b/updater_sample/tests/Android.mk
@@ -23,9 +23,9 @@
 LOCAL_JAVA_LIBRARIES := \
     android.test.base.stubs \
     android.test.runner.stubs \
-    guava \
+    guava
+LOCAL_STATIC_JAVA_LIBRARIES := android-support-test \
     mockito-target-minus-junit4
-LOCAL_STATIC_JAVA_LIBRARIES := android-support-test
 LOCAL_INSTRUMENTATION_FOR := SystemUpdaterSample
 LOCAL_PROGUARD_ENABLED := disabled
 
diff --git a/updater_sample/tests/res/raw/update_config_stream_001.json b/updater_sample/tests/res/raw/update_config_stream_001.json
index 15127cf..be51b7c 100644
--- a/updater_sample/tests/res/raw/update_config_stream_001.json
+++ b/updater_sample/tests/res/raw/update_config_stream_001.json
@@ -10,5 +10,8 @@
                 "size": 8
             }
         ]
+    },
+    "ab_config": {
+        "force_switch_slot": true
     }
 }
diff --git a/updater_sample/tests/res/raw/update_config_stream_002.json b/updater_sample/tests/res/raw/update_config_stream_002.json
index cf4469b..5d7874c 100644
--- a/updater_sample/tests/res/raw/update_config_stream_002.json
+++ b/updater_sample/tests/res/raw/update_config_stream_002.json
@@ -1,5 +1,8 @@
 {
     "__": "*** Generated using tools/gen_update_config.py ***",
+    "ab_config": {
+        "force_switch_slot": false
+    },
     "ab_install_type": "STREAMING",
     "ab_streaming_metadata": {
         "property_files": [
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
index 0975e76..000f566 100644
--- a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
+++ b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
@@ -18,6 +18,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
 
 import android.content.Context;
 import android.support.test.InstrumentationRegistry;
@@ -45,7 +46,8 @@
 
     private static final String JSON_NON_STREAMING =
             "{\"name\": \"vip update\", \"url\": \"file:///builds/a.zip\", "
-            + " \"ab_install_type\": \"NON_STREAMING\"}";
+                    + " \"ab_install_type\": \"NON_STREAMING\","
+                    + " \"ab_config\": { \"force_switch_slot\": false } }";
 
     @Rule
     public final ExpectedException thrown = ExpectedException.none();
@@ -82,6 +84,7 @@
                 config.getStreamingMetadata().getPropertyFiles()[0].getFilename());
         assertEquals(195, config.getStreamingMetadata().getPropertyFiles()[0].getOffset());
         assertEquals(8, config.getStreamingMetadata().getPropertyFiles()[0].getSize());
+        assertTrue(config.getAbConfig().getForceSwitchSlot());
     }
 
     @Test
@@ -94,7 +97,8 @@
     @Test
     public void getUpdatePackageFile_throwsErrorIfNotAFile() throws Exception {
         String json = "{\"name\": \"upd\", \"url\": \"http://foo.bar\","
-                + " \"ab_install_type\": \"NON_STREAMING\"}";
+                + " \"ab_install_type\": \"NON_STREAMING\","
+                + " \"ab_config\": { \"force_switch_slot\": false } }";
         UpdateConfig config = UpdateConfig.fromJson(json);
         thrown.expect(RuntimeException.class);
         config.getUpdatePackageFile();
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
new file mode 100644
index 0000000..0657a5e
--- /dev/null
+++ b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.android.systemupdatersample;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.os.UpdateEngine;
+import android.os.UpdateEngineCallback;
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
+
+import com.example.android.systemupdatersample.util.PayloadSpecs;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+import java.util.function.IntConsumer;
+
+/**
+ * Tests for {@link UpdateManager}
+ */
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class UpdateManagerTest {
+
+    @Rule
+    public MockitoRule mockito = MockitoJUnit.rule();
+
+    @Mock
+    private UpdateEngine mUpdateEngine;
+    @Mock
+    private PayloadSpecs mPayloadSpecs;
+    private UpdateManager mUpdateManager;
+
+    @Before
+    public void setUp() {
+        mUpdateManager = new UpdateManager(mUpdateEngine, mPayloadSpecs);
+    }
+
+    @Test
+    public void storesProgressThenInvokesCallbacks() {
+        IntConsumer statusUpdateCallback = mock(IntConsumer.class);
+
+        // When UpdateManager is bound to update_engine, it passes
+        // UpdateManager.UpdateEngineCallbackImpl as a callback to update_engine.
+        when(mUpdateEngine.bind(any(UpdateEngineCallback.class))).thenAnswer(answer -> {
+            UpdateEngineCallback callback = answer.getArgument(0);
+            callback.onStatusUpdate(/*engineStatus*/ 4, /*engineProgress*/ 0.2f);
+            return null;
+        });
+
+        mUpdateManager.setOnEngineStatusUpdateCallback(statusUpdateCallback);
+
+        // Making sure that manager.getProgress() returns correct progress
+        // in "onEngineStatusUpdate" callback.
+        doAnswer(answer -> {
+            assertEquals(0.2f, mUpdateManager.getProgress(), 1E-5);
+            return null;
+        }).when(statusUpdateCallback).accept(anyInt());
+
+        mUpdateManager.bind();
+
+        verify(statusUpdateCallback, times(1)).accept(4);
+    }
+
+}
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/ui/MainActivityTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/ui/MainActivityTest.java
deleted file mode 100644
index 0101416..0000000
--- a/updater_sample/tests/src/com/example/android/systemupdatersample/ui/MainActivityTest.java
+++ /dev/null
@@ -1,48 +0,0 @@
-/*
- * Copyright (C) 2018 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.example.android.systemupdatersample.ui;
-
-import static org.junit.Assert.assertNotNull;
-
-import android.support.test.filters.MediumTest;
-import android.support.test.rule.ActivityTestRule;
-import android.support.test.runner.AndroidJUnit4;
-
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-
-/**
- * Make sure that the main launcher activity opens up properly, which will be
- * verified by {@link #activityLaunches}.
- */
-@RunWith(AndroidJUnit4.class)
-@MediumTest
-public class MainActivityTest {
-
-    @Rule
-    public final ActivityTestRule<MainActivity> mActivityRule =
-            new ActivityTestRule<>(MainActivity.class);
-
-    /**
-     * Verifies that the activity under test can be launched.
-     */
-    @Test
-    public void activityLaunches() {
-        assertNotNull(mActivityRule.getActivity());
-    }
-}
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/util/PayloadSpecsTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/util/PayloadSpecsTest.java
index d9e5465..3ba84c1 100644
--- a/updater_sample/tests/src/com/example/android/systemupdatersample/util/PayloadSpecsTest.java
+++ b/updater_sample/tests/src/com/example/android/systemupdatersample/util/PayloadSpecsTest.java
@@ -55,6 +55,8 @@
     private Context mTargetContext;
     private Context mTestContext;
 
+    private PayloadSpecs mPayloadSpecs;
+
     @Rule
     public final ExpectedException thrown = ExpectedException.none();
 
@@ -64,6 +66,7 @@
         mTestContext = InstrumentationRegistry.getContext();
 
         mTestDir = mTargetContext.getFilesDir();
+        mPayloadSpecs = new PayloadSpecs();
     }
 
     @Test
@@ -75,7 +78,7 @@
         java.nio.file.Files.deleteIfExists(packageFile.toPath());
         java.nio.file.Files.copy(mTestContext.getResources().openRawResource(R.raw.ota_002_package),
                 packageFile.toPath());
-        PayloadSpec spec = PayloadSpecs.forNonStreaming(packageFile);
+        PayloadSpec spec = mPayloadSpecs.forNonStreaming(packageFile);
 
         assertEquals("correct url", "file://" + packageFile.getAbsolutePath(), spec.getUrl());
         assertEquals("correct payload offset",
@@ -90,7 +93,7 @@
     @Test
     public void forNonStreaming_IOException() throws Exception {
         thrown.expect(IOException.class);
-        PayloadSpecs.forNonStreaming(new File("/fake/news.zip"));
+        mPayloadSpecs.forNonStreaming(new File("/fake/news.zip"));
     }
 
     @Test
@@ -100,7 +103,7 @@
         long size = 200;
         File propertiesFile = createMockPropertiesFile();
 
-        PayloadSpec spec = PayloadSpecs.forStreaming(url, offset, size, propertiesFile);
+        PayloadSpec spec = mPayloadSpecs.forStreaming(url, offset, size, propertiesFile);
         assertEquals("same url", url, spec.getUrl());
         assertEquals("same offset", offset, spec.getOffset());
         assertEquals("same size", size, spec.getSize());
diff --git a/updater_sample/tools/gen_update_config.py b/updater_sample/tools/gen_update_config.py
index 4efa9f1..7fb64f7 100755
--- a/updater_sample/tools/gen_update_config.py
+++ b/updater_sample/tools/gen_update_config.py
@@ -46,10 +46,11 @@
     AB_INSTALL_TYPE_STREAMING = 'STREAMING'
     AB_INSTALL_TYPE_NON_STREAMING = 'NON_STREAMING'
 
-    def __init__(self, package, url, ab_install_type):
+    def __init__(self, package, url, ab_install_type, ab_force_switch_slot):
         self.package = package
         self.url = url
         self.ab_install_type = ab_install_type
+        self.ab_force_switch_slot = ab_force_switch_slot
         self.streaming_required = (
             # payload.bin and payload_properties.txt must exist.
             'payload.bin',
@@ -80,6 +81,9 @@
             'url': self.url,
             'ab_streaming_metadata': streaming_metadata,
             'ab_install_type': self.ab_install_type,
+            'ab_config': {
+                'force_switch_slot': self.ab_force_switch_slot,
+            }
         }
 
     def _gen_ab_streaming_metadata(self):
@@ -126,6 +130,11 @@
                         default=GenUpdateConfig.AB_INSTALL_TYPE_NON_STREAMING,
                         choices=ab_install_type_choices,
                         help='A/B update installation type')
+    parser.add_argument('--ab_force_switch_slot',
+                        type=bool,
+                        default=False,
+                        help='if set true device will boot to a new slot, otherwise user manually '
+                             'switches slot on the screen')
     parser.add_argument('package',
                         type=str,
                         help='OTA package zip file')
@@ -144,7 +153,8 @@
     gen = GenUpdateConfig(
         package=args.package,
         url=args.url,
-        ab_install_type=args.ab_install_type)
+        ab_install_type=args.ab_install_type,
+        ab_force_switch_slot=args.ab_force_switch_slot)
     gen.run()
     gen.write(args.out)
     print('Config is written to ' + args.out)