diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index 959d8a75a..e786bc9dc 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -213,7 +213,8 @@ class SnapshotManager final {
 
     // Perform first-stage mapping of snapshot targets. This replaces init's
     // call to CreateLogicalPartitions when snapshots are present.
-    bool CreateLogicalAndSnapshotPartitions(const std::string& super_device);
+    bool CreateLogicalAndSnapshotPartitions(const std::string& super_device,
+                                            const std::chrono::milliseconds& timeout_ms = {});
 
     // This method should be called preceding any wipe or flash of metadata or
     // userdata. It is only valid in recovery or fastbootd, and it ensures that
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 88d6b8d79..88731dfb6 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -1346,7 +1346,8 @@ bool SnapshotManager::NeedSnapshotsInFirstStageMount() {
     }
 }
 
-bool SnapshotManager::CreateLogicalAndSnapshotPartitions(const std::string& super_device) {
+bool SnapshotManager::CreateLogicalAndSnapshotPartitions(
+        const std::string& super_device, const std::chrono::milliseconds& timeout_ms) {
     LOG(INFO) << "Creating logical partitions with snapshots as needed";
 
     auto lock = LockExclusive();
@@ -1372,6 +1373,7 @@ bool SnapshotManager::CreateLogicalAndSnapshotPartitions(const std::string& supe
                 .metadata = metadata.get(),
                 .partition = &partition,
                 .partition_opener = &opener,
+                .timeout_ms = timeout_ms,
         };
         std::string ignore_path;
         if (!MapPartitionWithSnapshot(lock.get(), std::move(params), &ignore_path)) {
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 47ac474e2..7de37dbd9 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -999,7 +999,7 @@ TEST_F(SnapshotUpdateTest, FullUpdateFlow) {
     auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
     ASSERT_NE(init, nullptr);
     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
-    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", 1s));
 
     // Check that the target partitions have the same content.
     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
@@ -1127,7 +1127,7 @@ TEST_F(SnapshotUpdateTest, TestRollback) {
     auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
     ASSERT_NE(init, nullptr);
     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
-    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", 1s));
 
     // Check that the target partitions have the same content.
     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
@@ -1139,7 +1139,7 @@ TEST_F(SnapshotUpdateTest, TestRollback) {
     init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_a"));
     ASSERT_NE(init, nullptr);
     ASSERT_FALSE(init->NeedSnapshotsInFirstStageMount());
-    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", 1s));
 
     // Assert that the source partitions aren't affected.
     for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
@@ -1516,7 +1516,7 @@ TEST_F(SnapshotUpdateTest, Hashtree) {
     auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
     ASSERT_NE(init, nullptr);
     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
-    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", 1s));
 
     // Check that the target partition have the same content. Hashtree and FEC extents
     // should be accounted for.
@@ -1639,7 +1639,8 @@ TEST_P(FlashAfterUpdateTest, FlashSlotAfterUpdate) {
     // Simulate shutting down the device.
     ASSERT_TRUE(UnmapAll());
 
-    if (std::get<1>(GetParam()) /* merge */) {
+    bool after_merge = std::get<1>(GetParam());
+    if (after_merge) {
         ASSERT_TRUE(InitiateMerge("_b"));
         // Simulate shutting down the device after merge has initiated.
         ASSERT_TRUE(UnmapAll());
@@ -1688,21 +1689,11 @@ TEST_P(FlashAfterUpdateTest, FlashSlotAfterUpdate) {
     auto init = SnapshotManager::NewForFirstStageMount(
             new TestDeviceInfo(fake_super, flashed_slot_suffix));
     ASSERT_NE(init, nullptr);
-    if (init->NeedSnapshotsInFirstStageMount()) {
-        ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
-    } else {
-        for (const auto& name : {"sys", "vnd"}) {
-            ASSERT_TRUE(CreateLogicalPartition(
-                    CreateLogicalPartitionParams{
-                            .block_device = fake_super,
-                            .metadata_slot = flashed_slot,
-                            .partition_name = name + flashed_slot_suffix,
-                            .timeout_ms = 1s,
-                            .partition_opener = opener_.get(),
-                    },
-                    &path));
-        }
+
+    if (flashed_slot && after_merge) {
+        ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
     }
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", 1s));
 
     // Check that the target partitions have the same content.
     for (const auto& name : {"sys", "vnd"}) {