libsnapshot: tests for public APIs.

Test: libsnapshot_test

Change-Id: I411ae32e77914845ed4037d7e67620598f8218cf
diff --git a/fs_mgr/libsnapshot/digital_storage.h b/fs_mgr/libsnapshot/digital_storage.h
new file mode 100644
index 0000000..210298e
--- /dev/null
+++ b/fs_mgr/libsnapshot/digital_storage.h
@@ -0,0 +1,77 @@
+// 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 <stdint.h>
+#include <stdlib.h>
+
+namespace android {
+namespace digital_storage {
+
+template <size_t Power>
+struct Size {
+    static constexpr size_t power = Power;
+    constexpr Size(uint64_t count) : value_(count) {}
+
+    constexpr uint64_t bytes() const { return value_ << (Power * 10); }
+    constexpr uint64_t count() const { return value_; }
+    operator uint64_t() const { return bytes(); }
+
+  private:
+    uint64_t value_;
+};
+
+using B = Size<0>;
+using KiB = Size<1>;
+using MiB = Size<2>;
+using GiB = Size<3>;
+
+constexpr B operator""_B(unsigned long long v) {  // NOLINT
+    return B{v};
+}
+
+constexpr KiB operator""_KiB(unsigned long long v) {  // NOLINT
+    return KiB{v};
+}
+
+constexpr MiB operator""_MiB(unsigned long long v) {  // NOLINT
+    return MiB{v};
+}
+
+constexpr GiB operator""_GiB(unsigned long long v) {  // NOLINT
+    return GiB{v};
+}
+
+template <typename Dest, typename Src>
+constexpr Dest size_cast(Src src) {
+    if (Src::power < Dest::power) {
+        return Dest(src.count() >> ((Dest::power - Src::power) * 10));
+    }
+    if (Src::power > Dest::power) {
+        return Dest(src.count() << ((Src::power - Dest::power) * 10));
+    }
+    return Dest(src.count());
+}
+
+static_assert((1_B).bytes() == 1);
+static_assert((1_KiB).bytes() == 1 << 10);
+static_assert((1_MiB).bytes() == 1 << 20);
+static_assert((1_GiB).bytes() == 1 << 30);
+static_assert(size_cast<KiB>(1_B).count() == 0);
+static_assert(size_cast<KiB>(1024_B).count() == 1);
+static_assert(size_cast<KiB>(1_MiB).count() == 1024);
+
+}  // namespace digital_storage
+}  // namespace android
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index f4b5a13..6bc09a1 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -200,7 +200,9 @@
     FRIEND_TEST(SnapshotTest, Merge);
     FRIEND_TEST(SnapshotTest, MergeCannotRemoveCow);
     FRIEND_TEST(SnapshotTest, NoMergeBeforeReboot);
+    FRIEND_TEST(SnapshotUpdateTest, SnapshotStatusFileWithoutCow);
     friend class SnapshotTest;
+    friend class SnapshotUpdateTest;
     friend struct AutoDeleteCowImage;
     friend struct AutoDeleteSnapshot;
     friend struct PartitionCowCreator;
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 086ee95..876b8f8 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -33,6 +33,7 @@
 #include <liblp/builder.h>
 #include <liblp/mock_property_fetcher.h>
 
+#include "digital_storage.h"
 #include "test_helpers.h"
 #include "utility.h"
 
@@ -46,10 +47,12 @@
 using android::fs_mgr::BlockDeviceInfo;
 using android::fs_mgr::CreateLogicalPartitionParams;
 using android::fs_mgr::DestroyLogicalPartition;
+using android::fs_mgr::GetPartitionGroupName;
 using android::fs_mgr::GetPartitionName;
 using android::fs_mgr::MetadataBuilder;
 using namespace ::testing;
 using namespace android::fs_mgr::testing;
+using namespace android::digital_storage;
 using namespace std::chrono_literals;
 using namespace std::string_literals;
 
@@ -59,7 +62,7 @@
 TestDeviceInfo* test_device = nullptr;
 std::string fake_super;
 
-static constexpr uint64_t kSuperSize = 16 * 1024 * 1024;
+static constexpr uint64_t kSuperSize = (16_MiB).bytes();
 
 class SnapshotTest : public ::testing::Test {
   public:
@@ -106,7 +109,7 @@
         std::vector<std::string> snapshots = {"test-snapshot", "test_partition_a",
                                               "test_partition_b"};
         for (const auto& snapshot : snapshots) {
-            DeleteSnapshotDevice(snapshot);
+            ASSERT_TRUE(DeleteSnapshotDevice(snapshot));
             DeleteBackingImage(image_manager_, snapshot + "-cow-img");
 
             auto status_file = sm->GetSnapshotStatusFilePath(snapshot);
@@ -212,17 +215,23 @@
         return true;
     }
 
-    void DeleteSnapshotDevice(const std::string& snapshot) {
-        DeleteDevice(snapshot);
-        DeleteDevice(snapshot + "-inner");
-        DeleteDevice(snapshot + "-cow");
-        ASSERT_TRUE(image_manager_->UnmapImageIfExists(snapshot + "-cow-img"));
-        DeleteDevice(snapshot + "-base");
-    }
-    void DeleteDevice(const std::string& device) {
-        if (dm_.GetState(device) != DmDeviceState::INVALID) {
-            ASSERT_TRUE(dm_.DeleteDevice(device));
+    AssertionResult DeleteSnapshotDevice(const std::string& snapshot) {
+        AssertionResult res = AssertionSuccess();
+        if (!(res = DeleteDevice(snapshot))) return res;
+        if (!(res = DeleteDevice(snapshot + "-inner"))) return res;
+        if (!(res = DeleteDevice(snapshot + "-cow"))) return res;
+        if (!image_manager_->UnmapImageIfExists(snapshot + "-cow-img")) {
+            return AssertionFailure() << "Cannot unmap image " << snapshot << "-cow-img";
         }
+        if (!(res = DeleteDevice(snapshot + "-base"))) return res;
+        return AssertionSuccess();
+    }
+
+    AssertionResult DeleteDevice(const std::string& device) {
+        if (!dm_.DeleteDeviceIfExists(device)) {
+            return AssertionFailure() << "Can't delete " << device;
+        }
+        return AssertionSuccess();
     }
 
     AssertionResult CreateCowImage(const std::string& name) {
@@ -465,7 +474,7 @@
     // Wait 1s, otherwise DeleteSnapshotDevice may fail with EBUSY.
     sleep(1);
     // Forcefully delete the snapshot device, so it looks like we just rebooted.
-    DeleteSnapshotDevice("test-snapshot");
+    ASSERT_TRUE(DeleteSnapshotDevice("test-snapshot"));
 
     // Map snapshot should fail now, because we're in a merge-complete state.
     ASSERT_TRUE(AcquireLock());
@@ -602,7 +611,7 @@
 
     // Now, reflash super. Note that we haven't called ProcessUpdateState, so the
     // status is still Merging.
-    DeleteSnapshotDevice("test_partition_b");
+    ASSERT_TRUE(DeleteSnapshotDevice("test_partition_b"));
     ASSERT_TRUE(init->image_manager()->UnmapImageIfExists("test_partition_b-cow-img"));
     FormatFakeSuper();
     ASSERT_TRUE(CreatePartition("test_partition_b", kDeviceSize));
@@ -615,6 +624,419 @@
     ASSERT_EQ(sm->GetUpdateState(), UpdateState::None);
 }
 
+class SnapshotUpdateTest : public SnapshotTest {
+  public:
+    void SetUp() override {
+        SnapshotTest::SetUp();
+        Cleanup();
+
+        // Cleanup() changes slot suffix, so initialize it again.
+        test_device->set_slot_suffix("_a");
+
+        ON_CALL(*GetMockedPropertyFetcher(), GetBoolProperty("ro.virtual_ab.enabled", _))
+                .WillByDefault(Return(true));
+
+        opener = std::make_unique<TestPartitionOpener>(fake_super);
+
+        // Initialize source partition metadata. sys_b is similar to system_other.
+        // Not using full name "system", "vendor", "product" because these names collide with the
+        // mapped partitions on the running device.
+        src = MetadataBuilder::New(*opener, "super", 0);
+        ASSERT_NE(nullptr, src);
+        auto partition = src->AddPartition("sys_a", 0);
+        ASSERT_NE(nullptr, partition);
+        ASSERT_TRUE(src->ResizePartition(partition, 3_MiB));
+        partition = src->AddPartition("vnd_a", 0);
+        ASSERT_NE(nullptr, partition);
+        ASSERT_TRUE(src->ResizePartition(partition, 3_MiB));
+        partition = src->AddPartition("prd_a", 0);
+        ASSERT_NE(nullptr, partition);
+        ASSERT_TRUE(src->ResizePartition(partition, 3_MiB));
+        partition = src->AddPartition("sys_b", 0);
+        ASSERT_NE(nullptr, partition);
+        ASSERT_TRUE(src->ResizePartition(partition, 1_MiB));
+        auto metadata = src->Export();
+        ASSERT_NE(nullptr, metadata);
+        ASSERT_TRUE(UpdatePartitionTable(*opener, "super", *metadata.get(), 0));
+
+        // Map source partitions. Additionally, map sys_b to simulate system_other after flashing.
+        std::string path;
+        for (const auto& name : {"sys_a", "vnd_a", "prd_a", "sys_b"}) {
+            ASSERT_TRUE(CreateLogicalPartition(
+                    CreateLogicalPartitionParams{
+                            .block_device = fake_super,
+                            .metadata_slot = 0,
+                            .partition_name = name,
+                            .timeout_ms = 1s,
+                            .partition_opener = opener.get(),
+                    },
+                    &path));
+            ASSERT_TRUE(WriteRandomData(path));
+            auto hash = GetHash(path);
+            ASSERT_TRUE(hash.has_value());
+            hashes_[name] = *hash;
+        }
+    }
+    void TearDown() override {
+        Cleanup();
+        SnapshotTest::TearDown();
+    }
+    void Cleanup() {
+        if (!image_manager_) {
+            InitializeState();
+        }
+        for (const auto& suffix : {"_a", "_b"}) {
+            test_device->set_slot_suffix(suffix);
+            EXPECT_TRUE(sm->CancelUpdate()) << suffix;
+        }
+        EXPECT_TRUE(UnmapAll());
+    }
+
+    static AssertionResult ResizePartitions(
+            MetadataBuilder* builder,
+            const std::vector<std::pair<std::string, uint64_t>>& partition_sizes) {
+        for (auto&& [name, size] : partition_sizes) {
+            auto partition = builder->FindPartition(name);
+            if (!partition) {
+                return AssertionFailure() << "Cannot find partition in metadata " << name;
+            }
+            if (!builder->ResizePartition(partition, size)) {
+                return AssertionFailure()
+                       << "Cannot resize partition " << name << " to " << size << " bytes";
+            }
+        }
+        return AssertionSuccess();
+    }
+
+    AssertionResult IsPartitionUnchanged(const std::string& name) {
+        std::string path;
+        if (!dm_.GetDmDevicePathByName(name, &path)) {
+            return AssertionFailure() << "Path of " << name << " cannot be determined";
+        }
+        auto hash = GetHash(path);
+        if (!hash.has_value()) {
+            return AssertionFailure() << "Cannot read partition " << name << ": " << path;
+        }
+        if (hashes_[name] != *hash) {
+            return AssertionFailure() << "Content of " << name << " has changed after the merge";
+        }
+        return AssertionSuccess();
+    }
+
+    std::optional<uint64_t> GetSnapshotSize(const std::string& name) {
+        if (!AcquireLock()) {
+            return std::nullopt;
+        }
+        auto local_lock = std::move(lock_);
+
+        SnapshotManager::SnapshotStatus status;
+        if (!sm->ReadSnapshotStatus(local_lock.get(), name, &status)) {
+            return std::nullopt;
+        }
+        return status.snapshot_size;
+    }
+
+    AssertionResult UnmapAll() {
+        for (const auto& name : {"sys", "vnd", "prd"}) {
+            if (!dm_.DeleteDeviceIfExists(name + "_a"s)) {
+                return AssertionFailure() << "Cannot unmap " << name << "_a";
+            }
+            if (!DeleteSnapshotDevice(name + "_b"s)) {
+                return AssertionFailure() << "Cannot delete snapshot " << name << "_b";
+            }
+        }
+        return AssertionSuccess();
+    }
+
+    std::unique_ptr<TestPartitionOpener> opener;
+    std::unique_ptr<MetadataBuilder> src;
+    std::map<std::string, std::string> hashes_;
+};
+
+// Test full update flow executed by update_engine. Some partitions uses super empty space,
+// some uses images, and some uses both.
+// Also test UnmapUpdateSnapshot unmaps everything.
+// Also test first stage mount and merge after this.
+TEST_F(SnapshotUpdateTest, FullUpdateFlow) {
+    // OTA client calls CancelUpdate then BeginUpdate before doing anything.
+    ASSERT_TRUE(sm->CancelUpdate());
+    ASSERT_TRUE(sm->BeginUpdate());
+
+    // OTA client blindly unmaps all partitions that are possibly mapped.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(sm->UnmapUpdateSnapshot(name));
+    }
+
+    // OTA client adjusts the partition sizes before giving it to CreateUpdateSnapshots.
+    auto tgt = MetadataBuilder::NewForUpdate(*opener, "super", 0, 1);
+    ASSERT_NE(nullptr, tgt);
+    // clang-format off
+    ASSERT_TRUE(ResizePartitions(tgt.get(), {
+            {"sys_b", 4_MiB}, // grows
+            {"vnd_b", 4_MiB}, // grows
+            {"prd_b", 4_MiB}, // grows
+    }));
+    // clang-format on
+
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(tgt.get(), "_b", src.get(), "_a", {}));
+
+    // Test that partitions prioritize using space in super.
+    ASSERT_NE(nullptr, tgt->FindPartition("sys_b-cow"));
+    ASSERT_NE(nullptr, tgt->FindPartition("vnd_b-cow"));
+    ASSERT_EQ(nullptr, tgt->FindPartition("prd_b-cow"));
+
+    // The metadata slot 1 is now updated.
+    auto metadata = tgt->Export();
+    ASSERT_NE(nullptr, metadata);
+    ASSERT_TRUE(UpdatePartitionTable(*opener, "super", *metadata.get(), 1));
+
+    // Write some data to target partitions.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        std::string path;
+        ASSERT_TRUE(sm->MapUpdateSnapshot(
+                CreateLogicalPartitionParams{
+                        .block_device = fake_super,
+                        .metadata_slot = 1,
+                        .partition_name = name,
+                        .timeout_ms = 10s,
+                        .partition_opener = opener.get(),
+                },
+                &path))
+                << name;
+        ASSERT_TRUE(WriteRandomData(path));
+        auto hash = GetHash(path);
+        ASSERT_TRUE(hash.has_value());
+        hashes_[name] = *hash;
+    }
+
+    // Assert that source partitions aren't affected.
+    for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+
+    ASSERT_TRUE(sm->FinishedSnapshotWrites());
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto rebooted = new TestDeviceInfo(fake_super);
+    rebooted->set_slot_suffix("_b");
+    auto init = SnapshotManager::NewForFirstStageMount(rebooted);
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
+
+    // Check that the target partitions have the same content.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+
+    // Initiate the merge and wait for it to be completed.
+    ASSERT_TRUE(init->InitiateMerge());
+    ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());
+
+    // Check that the target partitions have the same content after the merge.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name))
+                << "Content of " << name << " changes after the merge";
+    }
+}
+
+// Test that if new system partitions uses empty space in super, that region is not snapshotted.
+TEST_F(SnapshotUpdateTest, DirectWriteEmptySpace) {
+    auto tgt = MetadataBuilder::NewForUpdate(*opener, "super", 0, 1);
+    ASSERT_NE(nullptr, tgt);
+    // clang-format off
+    ASSERT_TRUE(ResizePartitions(tgt.get(), {
+            {"sys_b", 4_MiB}, // grows
+            // vnd_b and prd_b are unchanged
+    }));
+    // clang-format on
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(tgt.get(), "_b", src.get(), "_a", {}));
+
+    ASSERT_EQ(3_MiB, GetSnapshotSize("sys_b").value_or(0));
+}
+
+// Test that if new system partitions uses space of old vendor partition, that region is
+// snapshotted.
+TEST_F(SnapshotUpdateTest, SnapshotOldPartitions) {
+    auto tgt = MetadataBuilder::NewForUpdate(*opener, "super", 0, 1);
+    ASSERT_NE(nullptr, tgt);
+    // clang-format off
+    ASSERT_TRUE(ResizePartitions(tgt.get(), {
+            {"vnd_b", 2_MiB}, // shrinks
+            {"sys_b", 4_MiB}, // grows
+            // prd_b is unchanged
+    }));
+    // clang-format on
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(tgt.get(), "_b", src.get(), "_a", {}));
+
+    ASSERT_EQ(4_MiB, GetSnapshotSize("sys_b").value_or(0));
+}
+
+// Test that even if there seem to be empty space in target metadata, COW partition won't take
+// it because they are used by old partitions.
+TEST_F(SnapshotUpdateTest, CowPartitionDoNotTakeOldPartitions) {
+    auto tgt = MetadataBuilder::NewForUpdate(*opener, "super", 0, 1);
+    ASSERT_NE(nullptr, tgt);
+    // clang-format off
+    ASSERT_TRUE(ResizePartitions(tgt.get(), {
+            {"sys_b", 2_MiB}, // shrinks
+            // vnd_b and prd_b are unchanged
+    }));
+    // clang-format on
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(tgt.get(), "_b", src.get(), "_a", {}));
+
+    auto metadata = tgt->Export();
+    ASSERT_NE(nullptr, metadata);
+    std::vector<std::string> written;
+    // Write random data to all COW partitions in super
+    for (auto p : metadata->partitions) {
+        if (GetPartitionGroupName(metadata->groups[p.group_index]) != kCowGroupName) {
+            continue;
+        }
+        std::string path;
+        ASSERT_TRUE(CreateLogicalPartition(
+                CreateLogicalPartitionParams{
+                        .block_device = fake_super,
+                        .metadata = metadata.get(),
+                        .partition = &p,
+                        .timeout_ms = 1s,
+                        .partition_opener = opener.get(),
+                },
+                &path));
+        ASSERT_TRUE(WriteRandomData(path));
+        written.push_back(GetPartitionName(p));
+    }
+    ASSERT_FALSE(written.empty())
+            << "No COW partitions are created even if there are empty space in super partition";
+
+    // Make sure source partitions aren't affected.
+    for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+}
+
+// Test that it crashes after creating snapshot status file but before creating COW image, then
+// calling CreateUpdateSnapshots again works.
+TEST_F(SnapshotUpdateTest, SnapshotStatusFileWithoutCow) {
+    // Write some trash snapshot files to simulate leftovers from previous runs.
+    {
+        ASSERT_TRUE(AcquireLock());
+        auto local_lock = std::move(lock_);
+        ASSERT_TRUE(sm->WriteSnapshotStatus(local_lock.get(), "sys_b",
+                                            SnapshotManager::SnapshotStatus{}));
+        ASSERT_TRUE(image_manager_->CreateBackingImage("sys_b-cow-img", 1_MiB,
+                                                       IImageManager::CREATE_IMAGE_DEFAULT));
+    }
+
+    // Redo the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->UnmapUpdateSnapshot("sys_b"));
+    auto tgt = MetadataBuilder::NewForUpdate(*opener, "super", 0, 1);
+    ASSERT_NE(nullptr, tgt);
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(tgt.get(), "_b", src.get(), "_a", {}));
+
+    // The metadata slot 1 is now updated.
+    auto metadata = tgt->Export();
+    ASSERT_NE(nullptr, metadata);
+    ASSERT_TRUE(UpdatePartitionTable(*opener, "super", *metadata.get(), 1));
+
+    // Check that target partitions can be mapped.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        std::string path;
+        EXPECT_TRUE(sm->MapUpdateSnapshot(
+                CreateLogicalPartitionParams{
+                        .block_device = fake_super,
+                        .metadata_slot = 1,
+                        .partition_name = name,
+                        .timeout_ms = 10s,
+                        .partition_opener = opener.get(),
+                },
+                &path))
+                << name;
+    }
+}
+
+// Test that the old partitions are not modified.
+TEST_F(SnapshotUpdateTest, TestRollback) {
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->UnmapUpdateSnapshot("sys_b"));
+    auto tgt = MetadataBuilder::NewForUpdate(*opener, "super", 0, 1);
+    ASSERT_NE(nullptr, tgt);
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(tgt.get(), "_b", src.get(), "_a", {}));
+
+    // The metadata slot 1 is now updated.
+    auto metadata = tgt->Export();
+    ASSERT_NE(nullptr, metadata);
+    ASSERT_TRUE(UpdatePartitionTable(*opener, "super", *metadata.get(), 1));
+
+    // Write some data to target partitions.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        std::string path;
+        ASSERT_TRUE(sm->MapUpdateSnapshot(
+                CreateLogicalPartitionParams{
+                        .block_device = fake_super,
+                        .metadata_slot = 1,
+                        .partition_name = name,
+                        .timeout_ms = 10s,
+                        .partition_opener = opener.get(),
+                },
+                &path))
+                << name;
+        ASSERT_TRUE(WriteRandomData(path));
+        auto hash = GetHash(path);
+        ASSERT_TRUE(hash.has_value());
+        hashes_[name] = *hash;
+    }
+
+    // Assert that source partitions aren't affected.
+    for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+
+    ASSERT_TRUE(sm->FinishedSnapshotWrites());
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto rebooted = new TestDeviceInfo(fake_super);
+    rebooted->set_slot_suffix("_b");
+    auto init = SnapshotManager::NewForFirstStageMount(rebooted);
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
+
+    // Check that the target partitions have the same content.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+
+    // Simulate shutting down the device again.
+    ASSERT_TRUE(UnmapAll());
+    rebooted = new TestDeviceInfo(fake_super);
+    rebooted->set_slot_suffix("_a");
+    init = SnapshotManager::NewForFirstStageMount(rebooted);
+    ASSERT_NE(init, nullptr);
+    ASSERT_FALSE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
+
+    // Assert that the source partitions aren't affected.
+    for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+}
+
+// Test that if an update is applied but not booted into, it can be canceled.
+TEST_F(SnapshotUpdateTest, CancelAfterApply) {
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->FinishedSnapshotWrites());
+    ASSERT_TRUE(sm->CancelUpdate());
+}
+
 }  // namespace snapshot
 }  // namespace android
 
@@ -656,6 +1078,7 @@
     }
 
     // Clean up previous run.
+    SnapshotUpdateTest().Cleanup();
     SnapshotTest().Cleanup();
 
     // Use a separate image manager for our fake super partition.
diff --git a/fs_mgr/libsnapshot/test_helpers.cpp b/fs_mgr/libsnapshot/test_helpers.cpp
index f67dd21..e12ec77 100644
--- a/fs_mgr/libsnapshot/test_helpers.cpp
+++ b/fs_mgr/libsnapshot/test_helpers.cpp
@@ -14,11 +14,18 @@
 
 #include "test_helpers.h"
 
+#include <android-base/file.h>
+#include <android-base/logging.h>
+#include <android-base/unique_fd.h>
 #include <gtest/gtest.h>
+#include <openssl/sha.h>
 
 namespace android {
 namespace snapshot {
 
+using android::base::ReadFully;
+using android::base::unique_fd;
+using android::base::WriteFully;
 using android::fiemap::IImageManager;
 
 void DeleteBackingImage(IImageManager* manager, const std::string& name) {
@@ -53,5 +60,55 @@
     return PartitionOpener::GetDeviceString(partition_name);
 }
 
+bool WriteRandomData(const std::string& path) {
+    unique_fd rand(open("/dev/urandom", O_RDONLY));
+    unique_fd fd(open(path.c_str(), O_WRONLY));
+
+    char buf[4096];
+    while (true) {
+        ssize_t n = TEMP_FAILURE_RETRY(read(rand.get(), buf, sizeof(buf)));
+        if (n <= 0) return false;
+        if (!WriteFully(fd.get(), buf, n)) {
+            if (errno == ENOSPC) {
+                return true;
+            }
+            PLOG(ERROR) << "Cannot write " << path;
+            return false;
+        }
+    }
+}
+
+std::string ToHexString(const uint8_t* buf, size_t len) {
+    char lookup[] = "0123456789abcdef";
+    std::string out(len * 2 + 1, '\0');
+    char* outp = out.data();
+    for (; len > 0; len--, buf++) {
+        *outp++ = (char)lookup[*buf >> 4];
+        *outp++ = (char)lookup[*buf & 0xf];
+    }
+    return out;
+}
+
+std::optional<std::string> GetHash(const std::string& path) {
+    unique_fd fd(open(path.c_str(), O_RDONLY));
+    char buf[4096];
+    SHA256_CTX ctx;
+    SHA256_Init(&ctx);
+    while (true) {
+        ssize_t n = TEMP_FAILURE_RETRY(read(fd.get(), buf, sizeof(buf)));
+        if (n < 0) {
+            PLOG(ERROR) << "Cannot read " << path;
+            return std::nullopt;
+        }
+        if (n == 0) {
+            break;
+        }
+        SHA256_Update(&ctx, buf, n);
+    }
+    uint8_t out[32];
+    SHA256_Final(out, &ctx);
+    return ToHexString(out, sizeof(out));
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/test_helpers.h b/fs_mgr/libsnapshot/test_helpers.h
index 19692ea..d917e35 100644
--- a/fs_mgr/libsnapshot/test_helpers.h
+++ b/fs_mgr/libsnapshot/test_helpers.h
@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <optional>
 #include <string>
 
 #include <libfiemap/image_manager.h>
@@ -66,5 +67,10 @@
 // Helper for error-spam-free cleanup.
 void DeleteBackingImage(android::fiemap::IImageManager* manager, const std::string& name);
 
+// Write some random data to the given device. Will write until reaching end of the device.
+bool WriteRandomData(const std::string& device);
+
+std::optional<std::string> GetHash(const std::string& path);
+
 }  // namespace snapshot
 }  // namespace android