diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 739fcbb..10cb7b8 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -1456,7 +1456,7 @@
                                             std::vector<std::string>* snapuserd_argv) {
     LOG(INFO) << "Performing transition for snapuserd.";
 
-    // Don't use EnsuerSnapuserdConnected() because this is called from init,
+    // Don't use EnsureSnapuserdConnected() because this is called from init,
     // and attempting to do so will deadlock.
     if (!snapuserd_client_ && transition != InitTransition::SELINUX_DETACH) {
         snapuserd_client_ = SnapuserdClient::Connect(kSnapuserdSocket, 10s);
@@ -1513,8 +1513,15 @@
             continue;
         }
 
+        std::string source_device_name;
+        if (snapshot_status.old_partition_size() > 0) {
+            source_device_name = GetSourceDeviceName(snapshot);
+        } else {
+            source_device_name = GetBaseDeviceName(snapshot);
+        }
+
         std::string source_device;
-        if (!dm.GetDmDevicePathByName(GetSourceDeviceName(snapshot), &source_device)) {
+        if (!dm.GetDmDevicePathByName(source_device_name, &source_device)) {
             LOG(ERROR) << "Could not get device path for " << GetSourceDeviceName(snapshot);
             continue;
         }
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 7da17bd..7630efe 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -2071,14 +2071,24 @@
     // After reboot, init does first stage mount.
     auto init = NewManagerForFirstStageMount("_b");
     ASSERT_NE(init, nullptr);
+
+    ASSERT_TRUE(init->EnsureSnapuserdConnected());
+    init->set_use_first_stage_snapuserd(true);
+
     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
 
     // Check that the target partitions have the same content.
-    for (const auto& name : {"sys_b", "vnd_b", "prd_b", "dlkm_b"}) {
+    std::vector<std::string> partitions = {"sys_b", "vnd_b", "prd_b", "dlkm_b"};
+    for (const auto& name : partitions) {
         ASSERT_TRUE(IsPartitionUnchanged(name));
     }
 
+    ASSERT_TRUE(init->PerformInitTransition(SnapshotManager::InitTransition::SECOND_STAGE));
+    for (const auto& name : partitions) {
+        ASSERT_TRUE(init->snapuserd_client()->WaitForDeviceDelete(name + "-user-cow-init"));
+    }
+
     // Initiate the merge and wait for it to be completed.
     ASSERT_TRUE(init->InitiateMerge());
     ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());