diff --git a/Android.bp b/Android.bp
index a4b7978..7cdd64c 100644
--- a/Android.bp
+++ b/Android.bp
@@ -763,6 +763,7 @@
 
     srcs: [
         "aosp/apex_handler_android_unittest.cc",
+        "aosp/cleanup_previous_update_action_unittest.cc",
         "aosp/dynamic_partition_control_android_unittest.cc",
         "aosp/update_attempter_android_unittest.cc",
         "certificate_checker_unittest.cc",
diff --git a/aosp/cleanup_previous_update_action_unittest.cc b/aosp/cleanup_previous_update_action_unittest.cc
new file mode 100644
index 0000000..0d2b4e6
--- /dev/null
+++ b/aosp/cleanup_previous_update_action_unittest.cc
@@ -0,0 +1,174 @@
+//
+// Copyright (C) 2021 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.
+//
+
+#include <algorithm>
+
+#include <brillo/message_loops/fake_message_loop.h>
+#include <gtest/gtest.h>
+#include <libsnapshot/snapshot.h>
+#include <libsnapshot/mock_snapshot.h>
+#include <libsnapshot/mock_snapshot_merge_stats.h>
+
+#include "update_engine/aosp/cleanup_previous_update_action.h"
+#include "update_engine/common/mock_boot_control.h"
+#include "update_engine/common/mock_dynamic_partition_control.h"
+#include "update_engine/common/mock_prefs.h"
+
+namespace chromeos_update_engine {
+
+using android::snapshot::AutoDevice;
+using android::snapshot::MockSnapshotManager;
+using android::snapshot::MockSnapshotMergeStats;
+using android::snapshot::UpdateState;
+using testing::_;
+using testing::AtLeast;
+using testing::Return;
+
+class MockCleanupPreviousUpdateActionDelegate final
+    : public CleanupPreviousUpdateActionDelegateInterface {
+  MOCK_METHOD(void, OnCleanupProgressUpdate, (double), (override));
+};
+
+class MockActionProcessor : public ActionProcessor {
+ public:
+  MOCK_METHOD(void, ActionComplete, (AbstractAction*, ErrorCode), (override));
+};
+
+class MockAutoDevice : public AutoDevice {
+ public:
+  explicit MockAutoDevice(std::string name) : AutoDevice(name) {}
+  ~MockAutoDevice() = default;
+};
+
+class CleanupPreviousUpdateActionTest : public ::testing::Test {
+ public:
+  void SetUp() override {
+    ON_CALL(boot_control_, GetDynamicPartitionControl())
+        .WillByDefault(Return(&dynamic_control_));
+    ON_CALL(boot_control_, GetCurrentSlot()).WillByDefault(Return(0));
+    ON_CALL(mock_snapshot_, GetSnapshotMergeStatsInstance())
+        .WillByDefault(Return(&mock_stats_));
+    action_.SetProcessor(&mock_processor_);
+    loop_.SetAsCurrent();
+  }
+
+  constexpr static FeatureFlag LAUNCH{FeatureFlag::Value::LAUNCH};
+  constexpr static FeatureFlag NONE{FeatureFlag::Value::NONE};
+  MockSnapshotManager mock_snapshot_;
+  MockPrefs mock_prefs_;
+  MockBootControl boot_control_;
+  MockDynamicPartitionControl dynamic_control_{};
+  MockCleanupPreviousUpdateActionDelegate mock_delegate_;
+  MockSnapshotMergeStats mock_stats_;
+  MockActionProcessor mock_processor_;
+  brillo::FakeMessageLoop loop_{nullptr};
+  CleanupPreviousUpdateAction action_{
+      &mock_prefs_, &boot_control_, &mock_snapshot_, &mock_delegate_};
+};
+
+TEST_F(CleanupPreviousUpdateActionTest, NonVabTest) {
+  // Since VAB isn't even enabled, |GetSnapshotMergeStatsInstance| shouldn't be
+  // called at all
+  EXPECT_CALL(mock_snapshot_, GetSnapshotMergeStatsInstance()).Times(0);
+  EXPECT_CALL(dynamic_control_, GetVirtualAbFeatureFlag())
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(NONE));
+  action_.PerformAction();
+}
+
+TEST_F(CleanupPreviousUpdateActionTest, VABSlotSuccessful) {
+  // Expectaion: if VABC is enabled, Clenup action should call
+  // |SnapshotMergeStats::Start()| to start merge, and wait for it to finish
+  EXPECT_CALL(mock_snapshot_, GetSnapshotMergeStatsInstance())
+      .Times(AtLeast(1));
+  EXPECT_CALL(mock_snapshot_, EnsureMetadataMounted())
+      .Times(AtLeast(1))
+      .WillRepeatedly(
+          []() { return std::make_unique<MockAutoDevice>("mock_device"); });
+  EXPECT_CALL(dynamic_control_, GetVirtualAbFeatureFlag())
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(LAUNCH));
+  // CleanupPreviousUpdateAction should use whatever slot returned by
+  // |GetCurrentSlot()|
+  EXPECT_CALL(boot_control_, GetCurrentSlot())
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(1));
+  EXPECT_CALL(boot_control_, IsSlotMarkedSuccessful(1))
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(true));
+  EXPECT_CALL(mock_snapshot_, ProcessUpdateState(_, _))
+      .Times(AtLeast(2))
+      .WillOnce(Return(UpdateState::Merging))
+      .WillRepeatedly(Return(UpdateState::MergeCompleted));
+  EXPECT_CALL(mock_stats_, Start())
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(true));
+  EXPECT_CALL(mock_processor_, ActionComplete(&action_, ErrorCode::kSuccess))
+      .Times(1);
+  action_.PerformAction();
+  while (loop_.PendingTasks()) {
+    ASSERT_TRUE(loop_.RunOnce(true));
+  }
+}
+
+TEST_F(CleanupPreviousUpdateActionTest, VabSlotNotReady) {
+  // Cleanup action should repeatly query boot control until the slot is marked
+  // successful.
+  static constexpr auto MAX_TIMEPOINT =
+      std::chrono::steady_clock::time_point::max();
+  EXPECT_CALL(mock_snapshot_, GetSnapshotMergeStatsInstance())
+      .Times(AtLeast(1));
+  EXPECT_CALL(mock_snapshot_, EnsureMetadataMounted())
+      .Times(AtLeast(1))
+      .WillRepeatedly(
+          []() { return std::make_unique<MockAutoDevice>("mock_device"); });
+  EXPECT_CALL(dynamic_control_, GetVirtualAbFeatureFlag())
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(LAUNCH));
+  auto slot_success_time = MAX_TIMEPOINT;
+  auto merge_start_time = MAX_TIMEPOINT;
+  EXPECT_CALL(boot_control_, IsSlotMarkedSuccessful(_))
+      .Times(AtLeast(3))
+      .WillOnce(Return(false))
+      .WillOnce(Return(false))
+      .WillOnce([&slot_success_time]() {
+        slot_success_time =
+            std::min(slot_success_time, std::chrono::steady_clock::now());
+        return true;
+      });
+
+  EXPECT_CALL(mock_stats_, Start())
+      .Times(1)
+      .WillRepeatedly([&merge_start_time]() {
+        merge_start_time =
+            std::min(merge_start_time, std::chrono::steady_clock::now());
+        return true;
+      });
+
+  EXPECT_CALL(mock_snapshot_, ProcessUpdateState(_, _))
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(UpdateState::MergeCompleted));
+  EXPECT_CALL(mock_processor_, ActionComplete(&action_, ErrorCode::kSuccess))
+      .Times(1);
+  action_.PerformAction();
+  while (loop_.PendingTasks()) {
+    ASSERT_TRUE(loop_.RunOnce(true));
+  }
+  ASSERT_LT(slot_success_time, merge_start_time)
+      << "Merge should not be started until slot is marked successful";
+}
+
+}  // namespace chromeos_update_engine
diff --git a/common/mock_boot_control.h b/common/mock_boot_control.h
new file mode 100644
index 0000000..f75ce5e
--- /dev/null
+++ b/common/mock_boot_control.h
@@ -0,0 +1,74 @@
+//
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#ifndef UPDATE_ENGINE_COMMON_MOCK_BOOT_CONTROL_H_
+#define UPDATE_ENGINE_COMMON_MOCK_BOOT_CONTROL_H_
+
+#include <memory>
+#include <string>
+
+#include <gmock/gmock.h>
+
+#include "update_engine/common/boot_control_stub.h"
+
+namespace chromeos_update_engine {
+
+class MockBootControl final : public BootControlStub {
+ public:
+  MOCK_METHOD(bool,
+              IsSlotMarkedSuccessful,
+              (BootControlInterface::Slot),
+              (const override));
+  MOCK_METHOD(unsigned int, GetNumSlots, (), (const override));
+  MOCK_METHOD(BootControlInterface::Slot, GetCurrentSlot, (), (const override));
+  MOCK_METHOD(bool,
+              GetPartitionDevice,
+              (const std::string&, Slot, bool, std::string*, bool*),
+              (const override));
+  MOCK_METHOD(bool,
+              GetPartitionDevice,
+              (const std::string&, BootControlInterface::Slot, std::string*),
+              (const override));
+  MOCK_METHOD(std::optional<PartitionDevice>,
+              GetPartitionDevice,
+              (const std::string&, uint32_t, uint32_t, bool),
+              (const override));
+
+  MOCK_METHOD(bool,
+              IsSlotBootable,
+              (BootControlInterface::Slot),
+              (const override));
+  MOCK_METHOD(bool,
+              MarkSlotUnbootable,
+              (BootControlInterface::Slot),
+              (override));
+  MOCK_METHOD(bool,
+              SetActiveBootSlot,
+              (BootControlInterface::Slot),
+              (override));
+  MOCK_METHOD(bool,
+              MarkBootSuccessfulAsync,
+              (base::Callback<void(bool)>),
+              (override));
+  MOCK_METHOD(DynamicPartitionControlInterface*,
+              GetDynamicPartitionControl,
+              (),
+              (override));
+};
+
+}  // namespace chromeos_update_engine
+
+#endif  // UPDATE_ENGINE_COMMON_MOCK_BOOT_CONTROL_H_