From 208957ba3c82c8faf33249b25252cddb3f3f72b1 Mon Sep 17 00:00:00 2001
From: Maksym Pavlenko <makpav@amazon.com>
Date: Fri, 22 Feb 2019 11:26:14 -0800
Subject: [PATCH] devmapper: proper cleanup in pool device test

Signed-off-by: Maksym Pavlenko <makpav@amazon.com>
---
 snapshots/devmapper/pool_device_test.go | 59 +++++++++++++------------
 1 file changed, 31 insertions(+), 28 deletions(-)

diff --git a/snapshots/devmapper/pool_device_test.go b/snapshots/devmapper/pool_device_test.go
index 9e423eff3..20ff57942 100644
--- a/snapshots/devmapper/pool_device_test.go
+++ b/snapshots/devmapper/pool_device_test.go
@@ -28,6 +28,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/containerd/containerd/mount"
 	"github.com/containerd/containerd/pkg/testutil"
 	"github.com/containerd/containerd/snapshots/devmapper/dmsetup"
 	"github.com/containerd/containerd/snapshots/devmapper/losetup"
@@ -107,37 +108,37 @@ func TestPoolDevice(t *testing.T) {
 	})
 
 	// Mount 'thin-1'
-	thin1MountPath := tempMountPath(t)
-	output, err := exec.Command("mount", dmsetup.GetFullDevicePath(thinDevice1), thin1MountPath).CombinedOutput()
-	assert.NilError(t, err, "failed to mount '%s': %s", thinDevice1, string(output))
+	err = mount.WithTempMount(ctx, getMounts(thinDevice1), func(thin1MountPath string) error {
+		// Write v1 test file on 'thin-1' device
+		thin1TestFilePath := filepath.Join(thin1MountPath, "TEST")
+		err := ioutil.WriteFile(thin1TestFilePath, []byte("test file (v1)"), 0700)
+		assert.NilError(t, err, "failed to write test file v1 on '%s' volume", thinDevice1)
 
-	// Write v1 test file on 'thin-1' device
-	thin1TestFilePath := filepath.Join(thin1MountPath, "TEST")
-	err = ioutil.WriteFile(thin1TestFilePath, []byte("test file (v1)"), 0700)
-	assert.NilError(t, err, "failed to write test file v1 on '%s' volume", thinDevice1)
+		// Take snapshot of 'thin-1'
+		t.Run("CreateSnapshotDevice", func(t *testing.T) {
+			testCreateSnapshot(t, pool)
+		})
 
-	// Take snapshot of 'thin-1'
-	t.Run("CreateSnapshotDevice", func(t *testing.T) {
-		testCreateSnapshot(t, pool)
+		// Update TEST file on 'thin-1' to v2
+		err = ioutil.WriteFile(thin1TestFilePath, []byte("test file (v2)"), 0700)
+		assert.NilError(t, err, "failed to write test file v2 on 'thin-1' volume after taking snapshot")
+
+		return nil
 	})
 
-	// Update TEST file on 'thin-1' to v2
-	err = ioutil.WriteFile(thin1TestFilePath, []byte("test file (v2)"), 0700)
-	assert.NilError(t, err, "failed to write test file v2 on 'thin-1' volume after taking snapshot")
+	assert.NilError(t, err)
 
 	// Mount 'snap-1' and make sure TEST file is v1
-	snap1MountPath := tempMountPath(t)
-	output, err = exec.Command("mount", dmsetup.GetFullDevicePath(snapDevice1), snap1MountPath).CombinedOutput()
-	assert.NilError(t, err, "failed to mount '%s' device: %s", snapDevice1, string(output))
+	err = mount.WithTempMount(ctx, getMounts(snapDevice1), func(snap1MountPath string) error {
+		// Read test file from snapshot device and make sure it's v1
+		fileData, err := ioutil.ReadFile(filepath.Join(snap1MountPath, "TEST"))
+		assert.NilError(t, err, "couldn't read test file from '%s' device", snapDevice1)
+		assert.Equal(t, "test file (v1)", string(fileData), "test file content is invalid on snapshot")
 
-	// Read test file from snapshot device and make sure it's v1
-	fileData, err := ioutil.ReadFile(filepath.Join(snap1MountPath, "TEST"))
-	assert.NilError(t, err, "couldn't read test file from '%s' device", snapDevice1)
-	assert.Assert(t, string(fileData) == "test file (v1)", "test file content is invalid on snapshot")
+		return nil
+	})
 
-	// Unmount devices before removing
-	output, err = exec.Command("umount", thin1MountPath, snap1MountPath).CombinedOutput()
-	assert.NilError(t, err, "failed to unmount devices: %s", string(output))
+	assert.NilError(t, err)
 
 	t.Run("DeactivateDevice", func(t *testing.T) {
 		testDeactivateThinDevice(t, pool)
@@ -207,11 +208,13 @@ func testRemoveThinDevice(t *testing.T, pool *PoolDevice) {
 	assert.NilError(t, err, "should delete thin device from pool")
 }
 
-func tempMountPath(t *testing.T) string {
-	path, err := ioutil.TempDir("", "devmapper-snapshotter-mount-")
-	assert.NilError(t, err, "failed to get temp directory for mount")
-
-	return path
+func getMounts(thinDeviceName string) []mount.Mount {
+	return []mount.Mount{
+		{
+			Source: dmsetup.GetFullDevicePath(thinDeviceName),
+			Type:   "ext4",
+		},
+	}
 }
 
 func createLoopbackDevice(t *testing.T, dir string) (string, string) {