Skip to content

Commit

Permalink
Rework TestSuccessfulClassicRun to improve cleaning
Browse files Browse the repository at this point in the history
  • Loading branch information
@upils
upils committed Oct 11, 2023
1 parent 46c1596 commit a780e32
Showing 1 changed file with 62 additions and 70 deletions.
132 changes: 62 additions & 70 deletions internal/statemachine/classic_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2167,13 +2167,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()
Expand All @@ -2187,9 +2187,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"),
Expand Down Expand Up @@ -2262,82 +2269,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())
}
}

Expand All @@ -2357,18 +2361,6 @@ 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())
}
}

err = stateMachine.Teardown()
asserter.AssertErrNil(err, true)

os.RemoveAll(stateMachine.stateMachineFlags.WorkDir)
})
}

Expand Down

0 comments on commit a780e32

Please sign in to comment.