diff --git a/fs_mgr/libsnapshot/Android.bp b/fs_mgr/libsnapshot/Android.bp
index 9256a16..1d72c70 100644
--- a/fs_mgr/libsnapshot/Android.bp
+++ b/fs_mgr/libsnapshot/Android.bp
@@ -72,6 +72,7 @@
     name: "libsnapshot_sources",
     srcs: [
         "android/snapshot/snapshot.proto",
+        "device_info.cpp",
         "snapshot.cpp",
         "snapshot_metadata_updater.cpp",
         "partition_cow_creator.cpp",
diff --git a/fs_mgr/libsnapshot/device_info.cpp b/fs_mgr/libsnapshot/device_info.cpp
new file mode 100644
index 0000000..bacb41c
--- /dev/null
+++ b/fs_mgr/libsnapshot/device_info.cpp
@@ -0,0 +1,123 @@
+// Copyright (C) 2019 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 "device_info.h"
+
+#include <android-base/logging.h>
+#include <fs_mgr.h>
+#include <fs_mgr_overlayfs.h>
+
+namespace android {
+namespace snapshot {
+
+#ifdef LIBSNAPSHOT_USE_HAL
+using android::hardware::boot::V1_0::CommandResult;
+#endif
+
+using namespace std::string_literals;
+
+#ifdef __ANDROID_RECOVERY__
+constexpr bool kIsRecovery = true;
+#else
+constexpr bool kIsRecovery = false;
+#endif
+
+std::string DeviceInfo::GetGsidDir() const {
+    return "ota"s;
+}
+
+std::string DeviceInfo::GetMetadataDir() const {
+    return "/metadata/ota"s;
+}
+
+std::string DeviceInfo::GetSlotSuffix() const {
+    return fs_mgr_get_slot_suffix();
+}
+
+std::string DeviceInfo::GetOtherSlotSuffix() const {
+    return fs_mgr_get_other_slot_suffix();
+}
+
+const android::fs_mgr::IPartitionOpener& DeviceInfo::GetPartitionOpener() const {
+    return opener_;
+}
+
+std::string DeviceInfo::GetSuperDevice(uint32_t slot) const {
+    return fs_mgr_get_super_partition_name(slot);
+}
+
+bool DeviceInfo::IsOverlayfsSetup() const {
+    return fs_mgr_overlayfs_is_setup();
+}
+
+#ifdef LIBSNAPSHOT_USE_HAL
+bool DeviceInfo::EnsureBootHal() {
+    if (!boot_control_) {
+        auto hal = android::hardware::boot::V1_0::IBootControl::getService();
+        if (!hal) {
+            LOG(ERROR) << "Could not find IBootControl HAL";
+            return false;
+        }
+        boot_control_ = android::hardware::boot::V1_1::IBootControl::castFrom(hal);
+        if (!boot_control_) {
+            LOG(ERROR) << "Could not find IBootControl 1.1 HAL";
+            return false;
+        }
+    }
+    return true;
+}
+#endif
+
+bool DeviceInfo::SetBootControlMergeStatus([[maybe_unused]] MergeStatus status) {
+#ifdef LIBSNAPSHOT_USE_HAL
+    if (!EnsureBootHal()) {
+        return false;
+    }
+    if (!boot_control_->setSnapshotMergeStatus(status)) {
+        LOG(ERROR) << "Unable to set the snapshot merge status";
+        return false;
+    }
+    return true;
+#else
+    LOG(ERROR) << "HAL support not enabled.";
+    return false;
+#endif
+}
+
+bool DeviceInfo::IsRecovery() const {
+    return kIsRecovery;
+}
+
+bool DeviceInfo::SetSlotAsUnbootable([[maybe_unused]] unsigned int slot) {
+#ifdef LIBSNAPSHOT_USE_HAL
+    if (!EnsureBootHal()) {
+        return false;
+    }
+
+    CommandResult result = {};
+    auto cb = [&](CommandResult r) -> void { result = r; };
+    boot_control_->setSlotAsUnbootable(slot, cb);
+    if (!result.success) {
+        LOG(ERROR) << "Error setting slot " << slot << " unbootable: " << result.errMsg;
+        return false;
+    }
+    return true;
+#else
+    LOG(ERROR) << "HAL support not enabled.";
+    return false;
+#endif
+}
+
+}  // namespace snapshot
+}  // namespace android
diff --git a/fs_mgr/libsnapshot/device_info.h b/fs_mgr/libsnapshot/device_info.h
new file mode 100644
index 0000000..d8d3d91
--- /dev/null
+++ b/fs_mgr/libsnapshot/device_info.h
@@ -0,0 +1,53 @@
+// Copyright (C) 2019 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.
+
+#pragma once
+
+#include <string>
+
+#ifdef LIBSNAPSHOT_USE_HAL
+#include <android/hardware/boot/1.1/IBootControl.h>
+#endif
+#include <liblp/partition_opener.h>
+#include <libsnapshot/snapshot.h>
+
+namespace android {
+namespace snapshot {
+
+class DeviceInfo final : public SnapshotManager::IDeviceInfo {
+    using MergeStatus = android::hardware::boot::V1_1::MergeStatus;
+
+  public:
+    std::string GetGsidDir() const override;
+    std::string GetMetadataDir() const override;
+    std::string GetSlotSuffix() const override;
+    std::string GetOtherSlotSuffix() const override;
+    const android::fs_mgr::IPartitionOpener& GetPartitionOpener() const override;
+    std::string GetSuperDevice(uint32_t slot) const override;
+    bool IsOverlayfsSetup() const override;
+    bool SetBootControlMergeStatus(MergeStatus status) override;
+    bool SetSlotAsUnbootable(unsigned int slot) override;
+    bool IsRecovery() const override;
+
+  private:
+    bool EnsureBootHal();
+
+    android::fs_mgr::PartitionOpener opener_;
+#ifdef LIBSNAPSHOT_USE_HAL
+    android::sp<android::hardware::boot::V1_1::IBootControl> boot_control_;
+#endif
+};
+
+}  // namespace snapshot
+}  // namespace android
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/auto_device.h b/fs_mgr/libsnapshot/include/libsnapshot/auto_device.h
index d5ceb0e..73450db 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/auto_device.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/auto_device.h
@@ -27,6 +27,8 @@
     virtual ~AutoDevice(){};
     void Release();
 
+    bool HasDevice() const { return !name_.empty(); }
+
   protected:
     AutoDevice(const std::string& name) : name_(name) {}
     std::string name_;
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index c820395..8f24709 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -122,6 +122,7 @@
         virtual const IPartitionOpener& GetPartitionOpener() const = 0;
         virtual bool IsOverlayfsSetup() const = 0;
         virtual bool SetBootControlMergeStatus(MergeStatus status) = 0;
+        virtual bool SetSlotAsUnbootable(unsigned int slot) = 0;
         virtual bool IsRecovery() const = 0;
     };
 
@@ -133,7 +134,7 @@
     static std::unique_ptr<SnapshotManager> New(IDeviceInfo* device = nullptr);
 
     // This is similar to New(), except designed specifically for first-stage
-    // init.
+    // init or recovery.
     static std::unique_ptr<SnapshotManager> NewForFirstStageMount(IDeviceInfo* device = nullptr);
 
     // Helper function for first-stage init to check whether a SnapshotManager
@@ -180,7 +181,10 @@
     //
     //   MergeCompleted indicates that the update has fully completed.
     //   GetUpdateState will return None, and a new update can begin.
-    UpdateState ProcessUpdateState();
+    //
+    // The optional callback allows the caller to periodically check the
+    // progress with GetUpdateState().
+    UpdateState ProcessUpdateState(const std::function<void()>& callback = {});
 
   public:
     // Initiate the merge if necessary, then wait for the merge to finish.
@@ -219,6 +223,18 @@
     // call to CreateLogicalPartitions when snapshots are present.
     bool CreateLogicalAndSnapshotPartitions(const std::string& super_device);
 
+    // This method should be called preceding any wipe or flash of metadata or
+    // userdata. It is only valid in recovery.
+    //
+    // When userdata will be wiped or flashed, it is necessary to clean up any
+    // snapshot state. If a merge is in progress, the merge must be finished.
+    // If a snapshot is present but not yet merged, the slot must be marked as
+    // unbootable.
+    //
+    // Returns true on success (or nothing to do), false on failure. The
+    // optional callback fires periodically to query progress via GetUpdateState.
+    bool HandleImminentDataWipe(const std::function<void()>& callback = {});
+
     // Dump debug information.
     bool Dump(std::ostream& os);
 
@@ -231,7 +247,7 @@
     // is not found in fstab.
     // Note: if this function is called the second time before the AutoDevice returned from the
     // first call is destroyed, the device will be unmounted when any of these AutoDevices is
-    // destroyed. FOr example:
+    // destroyed. For example:
     //   auto a = mgr->EnsureMetadataMounted(); // mounts
     //   auto b = mgr->EnsureMetadataMounted(); // does nothing
     //   b.reset() // unmounts
@@ -250,7 +266,10 @@
     FRIEND_TEST(SnapshotTest, Merge);
     FRIEND_TEST(SnapshotTest, NoMergeBeforeReboot);
     FRIEND_TEST(SnapshotTest, UpdateBootControlHal);
+    FRIEND_TEST(SnapshotUpdateTest, DataWipeAfterRollback);
+    FRIEND_TEST(SnapshotUpdateTest, DataWipeRollbackInRecovery);
     FRIEND_TEST(SnapshotUpdateTest, MergeCannotRemoveCow);
+    FRIEND_TEST(SnapshotUpdateTest, MergeInRecovery);
     FRIEND_TEST(SnapshotUpdateTest, SnapshotStatusFileWithoutCow);
     friend class SnapshotTest;
     friend class SnapshotUpdateTest;
@@ -464,6 +483,10 @@
             const LpMetadata* exported_target_metadata, const std::string& target_suffix,
             const std::map<std::string, SnapshotStatus>& all_snapshot_status);
 
+    // Unmap all partitions that were mapped by CreateLogicalAndSnapshotPartitions.
+    // This should only be called in recovery.
+    bool UnmapAllPartitions();
+
     std::string gsid_dir_;
     std::string metadata_dir_;
     std::unique_ptr<IDeviceInfo> device_;
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 2c516a2..1de7008 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -29,19 +29,16 @@
 #include <android-base/parseint.h>
 #include <android-base/strings.h>
 #include <android-base/unique_fd.h>
-#ifdef LIBSNAPSHOT_USE_HAL
-#include <android/hardware/boot/1.1/IBootControl.h>
-#endif
 #include <ext4_utils/ext4_utils.h>
 #include <fs_mgr.h>
 #include <fs_mgr_dm_linear.h>
-#include <fs_mgr_overlayfs.h>
 #include <fstab/fstab.h>
 #include <libdm/dm.h>
 #include <libfiemap/image_manager.h>
 #include <liblp/liblp.h>
 
 #include <android/snapshot/snapshot.pb.h>
+#include "device_info.h"
 #include "partition_cow_creator.h"
 #include "snapshot_metadata_updater.h"
 #include "utility.h"
@@ -77,57 +74,6 @@
 
 static constexpr char kBootIndicatorPath[] = "/metadata/ota/snapshot-boot";
 
-#ifdef __ANDROID_RECOVERY__
-constexpr bool kIsRecovery = true;
-#else
-constexpr bool kIsRecovery = false;
-#endif
-
-class DeviceInfo final : public SnapshotManager::IDeviceInfo {
-  public:
-    std::string GetGsidDir() const override { return "ota"s; }
-    std::string GetMetadataDir() const override { return "/metadata/ota"s; }
-    std::string GetSlotSuffix() const override { return fs_mgr_get_slot_suffix(); }
-    std::string GetOtherSlotSuffix() const override { return fs_mgr_get_other_slot_suffix(); }
-    const android::fs_mgr::IPartitionOpener& GetPartitionOpener() const { return opener_; }
-    std::string GetSuperDevice(uint32_t slot) const override {
-        return fs_mgr_get_super_partition_name(slot);
-    }
-    bool IsOverlayfsSetup() const override { return fs_mgr_overlayfs_is_setup(); }
-    bool SetBootControlMergeStatus(MergeStatus status) override;
-    bool IsRecovery() const override { return kIsRecovery; }
-
-  private:
-    android::fs_mgr::PartitionOpener opener_;
-#ifdef LIBSNAPSHOT_USE_HAL
-    android::sp<android::hardware::boot::V1_1::IBootControl> boot_control_;
-#endif
-};
-
-bool DeviceInfo::SetBootControlMergeStatus([[maybe_unused]] MergeStatus status) {
-#ifdef LIBSNAPSHOT_USE_HAL
-    if (!boot_control_) {
-        auto hal = android::hardware::boot::V1_0::IBootControl::getService();
-        if (!hal) {
-            LOG(ERROR) << "Could not find IBootControl HAL";
-            return false;
-        }
-        boot_control_ = android::hardware::boot::V1_1::IBootControl::castFrom(hal);
-        if (!boot_control_) {
-            LOG(ERROR) << "Could not find IBootControl 1.1 HAL";
-            return false;
-        }
-    }
-    if (!boot_control_->setSnapshotMergeStatus(status)) {
-        LOG(ERROR) << "Unable to set the snapshot merge status";
-        return false;
-    }
-    return true;
-#else
-    return false;
-#endif
-}
-
 // Note: IImageManager is an incomplete type in the header, so the default
 // destructor doesn't work.
 SnapshotManager::~SnapshotManager() {}
@@ -512,6 +458,14 @@
         return false;
     }
 
+    // We can't delete snapshots in recovery. The only way we'd try is it we're
+    // completing or canceling a merge in preparation for a data wipe, in which
+    // case, we don't care if the file sticks around.
+    if (device_->IsRecovery()) {
+        LOG(INFO) << "Skipping delete of snapshot " << name << " in recovery.";
+        return true;
+    }
+
     auto cow_image_name = GetCowImageDeviceName(name);
     if (images_->BackingImageExists(cow_image_name)) {
         if (!images_->DeleteBackingImage(cow_image_name)) {
@@ -744,7 +698,7 @@
 // Note that when a merge fails, we will *always* try again to complete the
 // merge each time the device boots. There is no harm in doing so, and if
 // the problem was transient, we might manage to get a new outcome.
-UpdateState SnapshotManager::ProcessUpdateState() {
+UpdateState SnapshotManager::ProcessUpdateState(const std::function<void()>& callback) {
     while (true) {
         UpdateState state = CheckMergeState();
         if (state == UpdateState::MergeFailed) {
@@ -756,6 +710,10 @@
             return state;
         }
 
+        if (callback) {
+            callback();
+        }
+
         // This wait is not super time sensitive, so we have a relatively
         // low polling frequency.
         std::this_thread::sleep_for(2s);
@@ -1718,9 +1676,9 @@
         case UpdateState::None:
         case UpdateState::MergeNeedsReboot:
         case UpdateState::MergeCompleted:
+        case UpdateState::Initiated:
             merge_status = MergeStatus::NONE;
             break;
-        case UpdateState::Initiated:
         case UpdateState::Unverified:
             merge_status = MergeStatus::SNAPSHOTTED;
             break;
@@ -2119,6 +2077,27 @@
     return UnmapPartitionWithSnapshot(lock.get(), target_partition_name);
 }
 
+bool SnapshotManager::UnmapAllPartitions() {
+    auto lock = LockExclusive();
+    if (!lock) return false;
+
+    const auto& opener = device_->GetPartitionOpener();
+    uint32_t slot = SlotNumberForSlotSuffix(device_->GetSlotSuffix());
+    auto super_device = device_->GetSuperDevice(slot);
+    auto metadata = android::fs_mgr::ReadMetadata(opener, super_device, slot);
+    if (!metadata) {
+        LOG(ERROR) << "Could not read dynamic partition metadata for device: " << super_device;
+        return false;
+    }
+
+    bool ok = true;
+    for (const auto& partition : metadata->partitions) {
+        auto partition_name = GetPartitionName(partition);
+        ok &= UnmapPartitionWithSnapshot(lock.get(), partition_name);
+    }
+    return ok;
+}
+
 bool SnapshotManager::Dump(std::ostream& os) {
     // Don't actually lock. Dump() is for debugging purposes only, so it is okay
     // if it is racy.
@@ -2192,5 +2171,66 @@
     return state;
 }
 
+bool SnapshotManager::HandleImminentDataWipe(const std::function<void()>& callback) {
+    if (!device_->IsRecovery()) {
+        LOG(ERROR) << "Data wipes are only allowed in recovery.";
+        return false;
+    }
+
+    auto mount = EnsureMetadataMounted();
+    if (!mount || !mount->HasDevice()) {
+        // We allow the wipe to continue, because if we can't mount /metadata,
+        // it is unlikely the device would have booted anyway. If there is no
+        // metadata partition, then the device predates Virtual A/B.
+        return true;
+    }
+
+    auto slot_number = SlotNumberForSlotSuffix(device_->GetSlotSuffix());
+    auto super_path = device_->GetSuperDevice(slot_number);
+    if (!CreateLogicalAndSnapshotPartitions(super_path)) {
+        LOG(ERROR) << "Unable to map partitions to complete merge.";
+        return false;
+    }
+
+    UpdateState state = ProcessUpdateState(callback);
+    LOG(INFO) << "Update state in recovery: " << state;
+    switch (state) {
+        case UpdateState::MergeFailed:
+            LOG(ERROR) << "Unrecoverable merge failure detected.";
+            return false;
+        case UpdateState::Unverified: {
+            // If an OTA was just applied but has not yet started merging, we
+            // have no choice but to revert slots, because the current slot will
+            // immediately become unbootable. Rather than wait for the device
+            // to reboot N times until a rollback, we proactively disable the
+            // new slot instead.
+            //
+            // Since the rollback is inevitable, we don't treat a HAL failure
+            // as an error here.
+            std::string old_slot;
+            auto boot_file = GetSnapshotBootIndicatorPath();
+            if (android::base::ReadFileToString(boot_file, &old_slot) &&
+                device_->GetSlotSuffix() != old_slot) {
+                LOG(ERROR) << "Reverting to slot " << old_slot << " since update will be deleted.";
+                device_->SetSlotAsUnbootable(slot_number);
+            }
+            break;
+        }
+        case UpdateState::MergeNeedsReboot:
+            // We shouldn't get here, because nothing is depending on
+            // logical partitions.
+            LOG(ERROR) << "Unexpected merge-needs-reboot state in recovery.";
+            break;
+        default:
+            break;
+    }
+
+    // Nothing should be depending on partitions now, so unmap them all.
+    if (!UnmapAllPartitions()) {
+        LOG(ERROR) << "Unable to unmap all partitions; fastboot may fail to flash.";
+    }
+    return true;
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 3c3d9a6..24354d1 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -1230,6 +1230,80 @@
     EXPECT_FALSE(IsMetadataMounted());
 }
 
+// Test that during a merge, we can wipe data in recovery.
+TEST_F(SnapshotUpdateTest, MergeInRecovery) {
+    // Execute the first update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(sm->FinishedSnapshotWrites());
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
+    init = nullptr;
+
+    // Initiate the merge and then immediately stop it to simulate a reboot.
+    auto new_sm = SnapshotManager::New(new TestDeviceInfo(fake_super, "_b"));
+    ASSERT_TRUE(new_sm->InitiateMerge());
+    ASSERT_TRUE(UnmapAll());
+
+    // Simulate a reboot into recovery.
+    auto test_device = std::make_unique<TestDeviceInfo>(fake_super, "_b");
+    test_device->set_recovery(true);
+    new_sm = SnapshotManager::NewForFirstStageMount(test_device.release());
+
+    ASSERT_TRUE(new_sm->HandleImminentDataWipe());
+    ASSERT_EQ(new_sm->GetUpdateState(), UpdateState::None);
+}
+
+// Test that after an OTA, before a merge, we can wipe data in recovery.
+TEST_F(SnapshotUpdateTest, DataWipeRollbackInRecovery) {
+    // Execute the first update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(sm->FinishedSnapshotWrites());
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // Simulate a reboot into recovery.
+    auto test_device = new TestDeviceInfo(fake_super, "_b");
+    test_device->set_recovery(true);
+    auto new_sm = SnapshotManager::NewForFirstStageMount(test_device);
+
+    ASSERT_TRUE(new_sm->HandleImminentDataWipe());
+    EXPECT_EQ(new_sm->GetUpdateState(), UpdateState::Unverified);
+    EXPECT_TRUE(test_device->IsSlotUnbootable(1));
+    EXPECT_FALSE(test_device->IsSlotUnbootable(0));
+}
+
+// Test that after an OTA and a bootloader rollback with no merge, we can wipe
+// data in recovery.
+TEST_F(SnapshotUpdateTest, DataWipeAfterRollback) {
+    // Execute the first update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(sm->FinishedSnapshotWrites());
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // Simulate a rollback, with reboot into recovery.
+    auto test_device = new TestDeviceInfo(fake_super, "_a");
+    test_device->set_recovery(true);
+    auto new_sm = SnapshotManager::NewForFirstStageMount(test_device);
+
+    ASSERT_TRUE(new_sm->HandleImminentDataWipe());
+    EXPECT_EQ(new_sm->GetUpdateState(), UpdateState::None);
+    EXPECT_FALSE(test_device->IsSlotUnbootable(0));
+    EXPECT_FALSE(test_device->IsSlotUnbootable(0));
+}
+
 class FlashAfterUpdateTest : public SnapshotUpdateTest,
                              public WithParamInterface<std::tuple<uint32_t, bool>> {
   public:
diff --git a/fs_mgr/libsnapshot/test_helpers.h b/fs_mgr/libsnapshot/test_helpers.h
index 0f70afe..9083843 100644
--- a/fs_mgr/libsnapshot/test_helpers.h
+++ b/fs_mgr/libsnapshot/test_helpers.h
@@ -16,6 +16,7 @@
 
 #include <optional>
 #include <string>
+#include <unordered_set>
 
 #include <android/hardware/boot/1.1/IBootControl.h>
 #include <gmock/gmock.h>
@@ -89,6 +90,12 @@
     }
     bool IsOverlayfsSetup() const override { return false; }
     bool IsRecovery() const override { return recovery_; }
+    bool SetSlotAsUnbootable(unsigned int slot) override {
+        unbootable_slots_.insert(slot);
+        return true;
+    }
+
+    bool IsSlotUnbootable(uint32_t slot) { return unbootable_slots_.count(slot) != 0; }
 
     void set_slot_suffix(const std::string& suffix) { slot_suffix_ = suffix; }
     void set_fake_super(const std::string& path) {
@@ -102,6 +109,7 @@
     std::unique_ptr<TestPartitionOpener> opener_;
     MergeStatus merge_status_;
     bool recovery_ = false;
+    std::unordered_set<uint32_t> unbootable_slots_;
 };
 
 class SnapshotTestPropertyFetcher : public android::fs_mgr::testing::MockPropertyFetcher {