Improve tests performance

internal/statemachine/classic_test.go

@@ -1887,6 +1887,8 @@ func TestFailedPopulateClassicRootfsContents(t *testing.T) {
 		err := stateMachine.makeTemporaryDirectories()
 		asserter.AssertErrNil(err, true)
 
+		t.Cleanup(func() { os.RemoveAll(stateMachine.stateMachineFlags.WorkDir) })
+
 		// also create chroot
 		err = stateMachine.createChroot()
 		asserter.AssertErrNil(err, true)
@@ -2167,13 +2169,13 @@ func TestFailedGenerateFilelist(t *testing.T) {
 func TestSuccessfulClassicRun(t *testing.T) {
 	t.Run("test_successful_classic_run", func(t *testing.T) {
 		asserter := helper.Asserter{T: t}
-		saveCWD := helper.SaveCWD()
-		defer saveCWD()
+		restoreCWD := helper.SaveCWD()
+		t.Cleanup(restoreCWD)
 
 		// We need the output directory set for this
 		outputDir, err := os.MkdirTemp("/tmp", "ubuntu-image-")
 		asserter.AssertErrNil(err, true)
-		defer os.RemoveAll(outputDir)
+		t.Cleanup(func() { os.RemoveAll(outputDir) })
 
 		var stateMachine ClassicStateMachine
 		stateMachine.commonFlags, stateMachine.stateMachineFlags = helper.InitCommonOpts()
@@ -2187,9 +2189,16 @@ func TestSuccessfulClassicRun(t *testing.T) {
 		err = stateMachine.Setup()
 		asserter.AssertErrNil(err, true)
 
+		t.Cleanup(func() { os.RemoveAll(stateMachine.stateMachineFlags.WorkDir) })
+
 		err = stateMachine.Run()
 		asserter.AssertErrNil(err, true)
 
+		t.Cleanup(func() {
+			err = stateMachine.Teardown()
+			asserter.AssertErrNil(err, true)
+		})
+
 		// make sure packages were successfully installed from public and private ppas
 		files := []string{
 			filepath.Join(stateMachine.tempDirs.chroot, "usr", "bin", "hello-ubuntu-image-public"),
@@ -2262,82 +2271,79 @@ func TestSuccessfulClassicRun(t *testing.T) {
 
 		// create a directory in which to mount the rootfs
 		mountDir := filepath.Join(stateMachine.tempDirs.scratch, "loopback")
-		// Slice used to store all the commands that need to be run
-		// to properly update grub.cfg in the chroot
 		var mountImageCmds []*exec.Cmd
 		var umountImageCmds []*exec.Cmd
 
+		t.Cleanup(func() {
+			for _, teardownCmd := range umountImageCmds {
+				if tmpErr := teardownCmd.Run(); tmpErr != nil {
+					if err != nil {
+						err = fmt.Errorf("%s after previous error: %w", tmpErr, err)
+					} else {
+						err = tmpErr
+					}
+				}
+			}
+		})
+
 		imgPath := filepath.Join(stateMachine.commonFlags.OutputDir, "pc-amd64.img")
 
 		// set up the loopback
 		mountImageCmds = append(mountImageCmds,
-			[]*exec.Cmd{
-				// set up the loopback
-				//nolint:gosec,G204
-				exec.Command("losetup",
-					filepath.Join("/dev", "loop99"),
-					imgPath,
-				),
-				//nolint:gosec,G204
-				exec.Command("kpartx",
-					"-a",
-					filepath.Join("/dev", "loop99"),
-				),
-				// mount the rootfs partition in which to run update-grub
-				//nolint:gosec,G204
-				exec.Command("mount",
-					filepath.Join("/dev", "mapper", "loop99p3"), // with this example the rootfs is partition 3
-					mountDir,
-				),
-			}...,
+			//nolint:gosec,G204
+			exec.Command("losetup",
+				filepath.Join("/dev", "loop99"),
+				imgPath,
+			),
+		)
+
+		// unset the loopback
+		umountImageCmds = append(umountImageCmds,
+			//nolint:gosec,G204
+			exec.Command("losetup", "--detach", filepath.Join("/dev", "loop99")),
+		)
+
+		mountImageCmds = append(mountImageCmds,
+			//nolint:gosec,G204
+			exec.Command("kpartx", "-a", filepath.Join("/dev", "loop99")),
+		)
+
+		umountImageCmds = append([]*exec.Cmd{
+			//nolint:gosec,G204
+			exec.Command("kpartx", "-d", filepath.Join("/dev", "loop99")),
+		}, umountImageCmds...,
+		)
+
+		mountImageCmds = append(mountImageCmds,
+			//nolint:gosec,G204
+			exec.Command("mount", filepath.Join("/dev", "mapper", "loop99p3"), mountDir), // with this example the rootfs is partition 3 mountDir
+		)
+
+		umountImageCmds = append([]*exec.Cmd{
+			//nolint:gosec,G204
+			exec.Command("mount", "--make-rprivate", filepath.Join("/dev", "mapper", "loop99p3")),
+			//nolint:gosec,G204
+			exec.Command("umount", "--recursive", filepath.Join("/dev", "mapper", "loop99p3")),
+		}, umountImageCmds...,
 		)
 
 		// set up the mountpoints
 		mountPoints := []string{"/dev", "/proc", "/sys"}
 		for _, mountPoint := range mountPoints {
 			mountCmds, umountCmds := mountFromHost(mountDir, mountPoint)
 			mountImageCmds = append(mountImageCmds, mountCmds...)
-			umountImageCmds = append(umountImageCmds, umountCmds...)
-			defer func(cmds []*exec.Cmd) {
-				_ = runAll(cmds)
-			}(umountCmds)
+			umountImageCmds = append(umountCmds, umountImageCmds...)
 		}
 		// make sure to unmount the disk too
-		umountImageCmds = append(umountImageCmds, exec.Command("umount", mountDir))
-
-		// tear down the loopback
-		teardownCmds := []*exec.Cmd{
-			//nolint:gosec,G204
-			exec.Command("kpartx",
-				"-d",
-				filepath.Join("/dev", "loop99"),
-			),
-			//nolint:gosec,G204
-			exec.Command("losetup",
-				"--detach",
-				filepath.Join("/dev", "loop99"),
-			),
-		}
-
-		for _, teardownCmd := range teardownCmds {
-			defer func(teardownCmd *exec.Cmd) {
-				if tmpErr := teardownCmd.Run(); tmpErr != nil {
-					if err != nil {
-						err = fmt.Errorf("%s after previous error: %w", tmpErr, err)
-					} else {
-						err = tmpErr
-					}
-				}
-
-			}(teardownCmd)
-		}
-		umountImageCmds = append(umountImageCmds, teardownCmds...)
+		umountImageCmds = append([]*exec.Cmd{exec.Command("umount", "--recursive", mountDir)}, umountImageCmds...)
 
 		// now run all the commands to mount the image
 		for _, cmd := range mountImageCmds {
+			outPut := helper.SetCommandOutput(cmd, true)
 			err := cmd.Run()
 			if err != nil {
-				t.Errorf("Error running command %s", cmd.String())
+				t.Errorf("Error running command \"%s\". Error is \"%s\". Output is: \n%s",
+					cmd.String(), err.Error(), outPut.String())
 			}
 		}
 
@@ -2357,18 +2363,64 @@ func TestSuccessfulClassicRun(t *testing.T) {
 			t.Errorf("Expected LANG=C.UTF-8 in %s, but got %s", localeFile, string(localeBytes))
 		}
 
-		// now run all the commands to unmount the image and clean up
-		for _, cmd := range umountImageCmds {
-			err := cmd.Run()
-			if err != nil {
-				t.Errorf("Error running command %s", cmd.String())
-			}
-		}
+	})
+}
+
+func TestSuccessfulRootfsGeneration(t *testing.T) {
+	t.Run("test_successful_classic_rootfs_generation", func(t *testing.T) {
+		asserter := helper.Asserter{T: t}
+		restoreCWD := helper.SaveCWD()
+		t.Cleanup(restoreCWD)
+
+		// We need the output directory set for this
+		outputDir, err := os.MkdirTemp("/tmp", "ubuntu-image-")
+		asserter.AssertErrNil(err, true)
+		t.Cleanup(func() { os.RemoveAll(outputDir) })
 
-		err = stateMachine.Teardown()
+		var stateMachine ClassicStateMachine
+		stateMachine.commonFlags, stateMachine.stateMachineFlags = helper.InitCommonOpts()
+		stateMachine.parent = &stateMachine
+		stateMachine.commonFlags.Debug = true
+		stateMachine.commonFlags.Size = "5G"
+		stateMachine.commonFlags.OutputDir = outputDir
+		stateMachine.Args.ImageDefinition = filepath.Join("testdata", "image_definitions",
+			"test_rootfs_tarball.yaml")
+
+		err = stateMachine.Setup()
 		asserter.AssertErrNil(err, true)
 
-		os.RemoveAll(stateMachine.stateMachineFlags.WorkDir)
+		t.Cleanup(func() { os.RemoveAll(stateMachine.stateMachineFlags.WorkDir) })
+
+		err = stateMachine.Run()
+		asserter.AssertErrNil(err, true)
+
+		t.Cleanup(func() {
+			err = stateMachine.Teardown()
+			asserter.AssertErrNil(err, true)
+		})
+
+		// make sure all the artifacts were created and are the correct file types
+		artifacts := map[string]string{
+			"rootfs.tar": "tar archive",
+		}
+		for artifact, fileType := range artifacts {
+			fullPath := filepath.Join(stateMachine.commonFlags.OutputDir, artifact)
+			_, err := os.Stat(fullPath)
+			if err != nil {
+				if os.IsNotExist(err) {
+					t.Errorf("File \"%s\" should exist, but does not", fullPath)
+				}
+			}
+
+			// check it is the expected file type
+			fileCommand := *exec.Command("file", fullPath)
+			cmdOutput, err := fileCommand.CombinedOutput()
+			asserter.AssertErrNil(err, true)
+			if !strings.Contains(string(cmdOutput), fileType) {
+				t.Errorf("File \"%s\" is the wrong file type. Expected \"%s\" but got \"%s\"",
+					fullPath, fileType, string(cmdOutput))
+			}
+		}
 	})
 }
 
@@ -3657,6 +3709,8 @@ func TestPreseedResetChroot(t *testing.T) {
 		err := stateMachine.makeTemporaryDirectories()
 		asserter.AssertErrNil(err, true)
 
+		t.Cleanup(func() { os.RemoveAll(stateMachine.stateMachineFlags.WorkDir) })
+
 		// create chroot to preseed
 		err = stateMachine.createChroot()
 		asserter.AssertErrNil(err, true)
@@ -3786,6 +3840,8 @@ func TestUnsupportedBootloader(t *testing.T) {
 		err := stateMachine.makeTemporaryDirectories()
 		asserter.AssertErrNil(err, true)
 
+		t.Cleanup(func() { os.RemoveAll(stateMachine.stateMachineFlags.WorkDir) })
+
 		// place a test gadget tree in the scratch directory so we don't
 		// have to build one
 		gadgetSource := filepath.Join("testdata", "gadget_tree")

internal/statemachine/helper.go

@@ -608,6 +608,7 @@ func cloneGitRepo(imageDefinition imagedefinition.ImageDefinition, workDir strin
 	cloneOptions := &git.CloneOptions{
 		URL:          imageDefinition.Gadget.GadgetURL,
 		SingleBranch: true,
+		Depth:        1,
 	}
 	if imageDefinition.Gadget.GadgetBranch != "" {
 		cloneOptions.ReferenceName = plumbing.NewBranchReferenceName(imageDefinition.Gadget.GadgetBranch)

internal/statemachine/state_machine.go

@@ -14,6 +14,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"time"
 
 	diskfs "github.com/diskfs/go-diskfs"
 	"github.com/snapcore/snapd/gadget"
@@ -441,7 +442,12 @@ func (stateMachine *StateMachine) Run() error {
 		if !stateMachine.commonFlags.Quiet {
 			fmt.Printf("[%d] %s\n", stateMachine.StepsTaken, stateFunc.name)
 		}
-		if err := stateFunc.function(stateMachine); err != nil {
+		start := time.Now()
+		err := stateFunc.function(stateMachine)
+		if stateMachine.commonFlags.Debug {
+			fmt.Printf("duration: %v\n", time.Since(start))
+		}
+		if err != nil {
 			// clean up work dir on error
 			cleanupErr := stateMachine.cleanup()
 			if cleanupErr != nil {

internal/statemachine/testdata/image_definitions/test_amd64.yaml

@@ -64,7 +64,5 @@ artifacts:
       name: pc-amd64.qcow2
   manifest:
     name: "filesystem-manifest.txt"
-  rootfs-tarball:
-    name: "rootfs.tar"
   filelist:
     name: "filesystem-filelist.txt"

internal/statemachine/testdata/image_definitions/test_rootfs_tarball.yaml

@@ -0,0 +1,32 @@
+name: ubuntu-server-amd64
+display-name: Ubuntu Server amd64
+revision: 1
+architecture: amd64
+series: jammy
+class: preinstalled
+kernel: linux-image-generic
+gadget:
+  url: "https://github.com/snapcore/pc-gadget.git"
+  branch: classic
+  type: "git"
+rootfs:
+  components:
+    - main
+    - universe
+    - restricted
+  seed:
+    urls:
+      - "git://git.launchpad.net/~ubuntu-core-dev/ubuntu-seeds/+git/"
+    branch: jammy
+    names:
+      - server
+      - minimal
+      - standard
+      - cloud-image
+customization:
+  extra-snaps:
+    - name: core
+    - name: core20
+artifacts:
+  rootfs-tarball:
+    name: "rootfs.tar"