Fix and add some tests cases

internal/statemachine/classic_states.go

@@ -541,23 +541,26 @@ func (stateMachine *StateMachine) createChroot() error {
 }
 
 // add PPAs to the apt sources list
-func (stateMachine *StateMachine) addExtraPPAs() error {
+func (stateMachine *StateMachine) addExtraPPAs() (err error) {
 	classicStateMachine := stateMachine.parent.(*ClassicStateMachine)
 
 	// create /etc/apt/sources.list.d in the chroot if it doesn't already exist
 	sourcesListD := filepath.Join(stateMachine.tempDirs.chroot, "etc", "apt", "sources.list.d")
-	err := osMkdir(sourcesListD, 0755)
+	err = osMkdir(sourcesListD, 0755)
 	if err != nil && !os.IsExist(err) {
-		return fmt.Errorf("Failed to create apt sources.list.d: %s", err.Error())
+		err = fmt.Errorf("Failed to create apt sources.list.d: %s", err.Error())
+		return err
 	}
 
 	// now create the ppa sources.list files
 	tmpGPGDir, err := osMkdirTemp("/tmp", "ubuntu-image-gpg")
 	if err != nil {
-		return fmt.Errorf("Error creating temp dir for gpg imports: %s", err.Error())
+		err = fmt.Errorf("Error creating temp dir for gpg imports: %s", err.Error())
+		return err
 	}
 	defer func() {
-		if tmpErr := osRemoveAll(tmpGPGDir); tmpErr != nil {
+		tmpErr := osRemoveAll(tmpGPGDir)
+		if tmpErr != nil {
 			if err != nil {
 				err = fmt.Errorf("%s after previous error: %w", tmpErr.Error(), err)
 			} else {
@@ -569,14 +572,17 @@ func (stateMachine *StateMachine) addExtraPPAs() error {
 		ppaFileName, ppaFileContents := createPPAInfo(ppa,
 			classicStateMachine.ImageDef.Series)
 
+		var ppaIO *os.File
 		ppaFile := filepath.Join(sourcesListD, ppaFileName)
-		ppaIO, err := osOpenFile(ppaFile, os.O_CREATE|os.O_WRONLY, 0644)
+		ppaIO, err = osOpenFile(ppaFile, os.O_CREATE|os.O_WRONLY, 0644)
 		if err != nil {
-			return fmt.Errorf("Error creating %s: %s", ppaFile, err.Error())
+			err = fmt.Errorf("Error creating %s: %s", ppaFile, err.Error())
+			return err
 		}
 		_, err = ppaIO.Write([]byte(ppaFileContents))
 		if err != nil {
-			return fmt.Errorf("unable to write ppa file %s: %w", ppaFile, err)
+			err = fmt.Errorf("unable to write ppa file %s: %w", ppaFile, err)
+			return err
 		}
 		ppaIO.Close()
 
@@ -591,12 +597,14 @@ func (stateMachine *StateMachine) addExtraPPAs() error {
 			"etc", "apt", "trusted.gpg.d", keyFileName)
 		err = importPPAKeys(ppa, tmpGPGDir, keyFilePath, stateMachine.commonFlags.Debug)
 		if err != nil {
-			return fmt.Errorf("Error retrieving signing key for ppa \"%s\": %s",
+			err = fmt.Errorf("Error retrieving signing key for ppa \"%s\": %s",
 				ppa.PPAName, err.Error())
+			return err
 		}
 	}
-	if err := osRemoveAll(tmpGPGDir); err != nil {
-		return fmt.Errorf("Error removing temporary gpg directory \"%s\": %s", tmpGPGDir, err.Error())
+	if err = osRemoveAll(tmpGPGDir); err != nil {
+		err = fmt.Errorf("Error removing temporary gpg directory \"%s\": %s", tmpGPGDir, err.Error())
+		return err
 	}
 
 	return nil
@@ -983,10 +991,8 @@ func (stateMachine *StateMachine) customizeFstab() error {
 		fstabEntries = append(fstabEntries, fstabEntry)
 	}
 	_, err = fstabIO.Write([]byte(strings.Join(fstabEntries, "\n") + "\n"))
-	if err != nil {
-		return err
-	}
-	return nil
+
+	return err
 }
 
 // Handle any manual customizations specified in the image definition

internal/statemachine/classic_test.go

@@ -2943,6 +2943,23 @@ func TestFailedAddExtraPPAs(t *testing.T) {
 		asserter.AssertErrContains(err, "Error removing temporary gpg directory")
 		osRemoveAll = os.RemoveAll
 
+		// Test failing osRemoveAll in defered function
+		// mock os.RemoveAll
+		osRemoveAll = mockRemoveAll
+		defer func() {
+			osRemoveAll = os.RemoveAll
+		}()
+		// mock os.OpenFile
+		osOpenFile = mockOpenFile
+		defer func() {
+			osOpenFile = os.OpenFile
+		}()
+		err = stateMachine.addExtraPPAs()
+		asserter.AssertErrContains(err, "Error creating")
+		asserter.AssertErrContains(err, "after previous error")
+		osRemoveAll = os.RemoveAll
+		osOpenFile = os.OpenFile
+
 		os.RemoveAll(stateMachine.stateMachineFlags.WorkDir)
 	})
 }

internal/statemachine/helper.go

@@ -973,10 +973,10 @@ func runAll(cmds []*exec.Cmd) error {
 }
 
 // updateGrub mounts the resulting image and runs update-grub
-func (stateMachine *StateMachine) updateGrub(rootfsVolName string, rootfsPartNum int) error {
+func (stateMachine *StateMachine) updateGrub(rootfsVolName string, rootfsPartNum int) (err error) {
 	// create a directory in which to mount the rootfs
 	mountDir := filepath.Join(stateMachine.tempDirs.scratch, "loopback")
-	err := osMkdir(mountDir, 0755)
+	err = osMkdir(mountDir, 0755)
 	if err != nil && !os.IsExist(err) {
 		return fmt.Errorf("Error creating scratch/loopback directory: %s", err.Error())
 	}
@@ -996,12 +996,14 @@ func (stateMachine *StateMachine) updateGrub(rootfsVolName string, rootfsPartNum
 		stateMachine.commonFlags.SectorSize,
 		imgPath,
 	)
-	losetupOutput, err := losetupCmd.Output()
+	var losetupOutput []byte
+	losetupOutput, err = losetupCmd.Output()
 	if err != nil {
-		return fmt.Errorf("Error running losetup command \"%s\". Error is %s",
+		err = fmt.Errorf("Error running losetup command \"%s\". Error is %s",
 			losetupCmd.String(),
 			err.Error(),
 		)
+		return err
 	}
 	loopUsed := strings.TrimSpace(string(losetupOutput))
 
@@ -1024,7 +1026,7 @@ func (stateMachine *StateMachine) updateGrub(rootfsVolName string, rootfsPartNum
 		defer func(cmds []*exec.Cmd) {
 			if tmpErr := runAll(cmds); tmpErr != nil {
 				if err != nil {
-					err = fmt.Errorf("%s after previous error: %w", tmpErr, err)
+					err = fmt.Errorf("Unable to unmount: %s after previous error: %w", tmpErr, err)
 				} else {
 					err = tmpErr
 				}
@@ -1082,7 +1084,7 @@ func (stateMachine *StateMachine) updateGrub(rootfsVolName string, rootfsPartNum
 	defer func() {
 		if tmpErr := teardownCmd.Run(); tmpErr != nil {
 			if err != nil {
-				err = fmt.Errorf("%s after previous error: %w", tmpErr, err)
+				err = fmt.Errorf("Unable to execute teardown cmd: %s after previous error: %w", tmpErr, err)
 			} else {
 				err = tmpErr
 			}
@@ -1093,10 +1095,11 @@ func (stateMachine *StateMachine) updateGrub(rootfsVolName string, rootfsPartNum
 	// now run all the commands
 	for _, cmd := range updateGrubCmds {
 		cmdOutput := helper.SetCommandOutput(cmd, stateMachine.commonFlags.Debug)
-		err := cmd.Run()
+		err = cmd.Run()
 		if err != nil {
-			return fmt.Errorf("Error running command \"%s\". Error is \"%s\". Output is: \n%s",
+			err = fmt.Errorf("Error running command \"%s\". Error is \"%s\". Output is: \n%s",
 				cmd.String(), err.Error(), cmdOutput.String())
+			return err
 		}
 	}
 

internal/statemachine/helper_test.go

@@ -1375,6 +1375,9 @@ func TestFailedUpdateGrub(t *testing.T) {
 		testCaseName = "TestFailedUpdateGrubOther"
 		err = stateMachine.updateGrub("", 0)
 		asserter.AssertErrContains(err, "Error running command")
+		// check defered function failed and wrapped the error
+		asserter.AssertErrContains(err, "Unable to execute teardown cmd")
+		asserter.AssertErrContains(err, "Unable to unmount")
 		execCommand = exec.Command
 	})
 }

internal/statemachine/snap_states.go

@@ -37,7 +37,7 @@ func (stateMachine *StateMachine) prepareImage() error {
 			fmt.Printf("WARNING: revision %d for snap %s may not be the latest available version!\n", snapRev, snapName)
 			err = imageOpts.SeedManifest.SetAllowedSnapRevision(snapName, snap.R(snapRev))
 			if err != nil {
-				return fmt.Errorf("error dealing with the extra snap %s: %w", snapName, err)
+				return fmt.Errorf("Error preparing image: error dealing with snap revision %s: %w", snapName, err)
 			}
 		}
 	}

internal/statemachine/snap_test.go

@@ -177,10 +177,11 @@ func TestSuccessfulSnapCore18(t *testing.T) {
 	})
 }
 
-// TestFailedPrepareImage tests a failure in the call to image.Prepare. This is easy to achieve
-// by attempting to use --disable-console-conf with a core20 image
+// TestFailedPrepareImage tests prepareImage
 func TestFailedPrepareImage(t *testing.T) {
-	t.Run("test_failed_prepare_image", func(t *testing.T) {
+	// Test a failure in the call to image.Prepare. This is easy to achieve
+	// by attempting to use --disable-console-conf with a core20 image
+	t.Run("test_failed_prepare_image_imagePrepare", func(t *testing.T) {
 		asserter := helper.Asserter{T: t}
 		saveCWD := helper.SaveCWD()
 		defer saveCWD()
@@ -195,11 +196,37 @@ func TestFailedPrepareImage(t *testing.T) {
 		asserter.AssertErrNil(err, true)
 
 		err = stateMachine.Run()
+		fmt.Print(err)
 		asserter.AssertErrContains(err, "Error preparing image")
 
 		err = stateMachine.Teardown()
 		asserter.AssertErrNil(err, true)
 	})
+
+	t.Run("test_failed_prepare_image_snap_revision", func(t *testing.T) {
+		asserter := helper.Asserter{T: t}
+		saveCWD := helper.SaveCWD()
+		defer saveCWD()
+
+		var stateMachine SnapStateMachine
+		stateMachine.commonFlags, stateMachine.stateMachineFlags = helper.InitCommonOpts()
+		stateMachine.parent = &stateMachine
+		stateMachine.Args.ModelAssertion = filepath.Join("testdata", "modelAssertion20")
+		stateMachine.Opts.Revisions = map[string]int{
+			"test": 0,
+		}
+
+		err := stateMachine.Setup()
+		asserter.AssertErrNil(err, true)
+
+		err = stateMachine.Run()
+		fmt.Print(err)
+		asserter.AssertErrContains(err, "error dealing with snap revision")
+
+		err = stateMachine.Teardown()
+		asserter.AssertErrNil(err, true)
+	})
+
 }
 
 // TestPopulateSnapRootfsContents runs the state machine through populate_rootfs_contents and examines