From 862f1117c0f12a4c8a612bcb393305ceed46cc96 Mon Sep 17 00:00:00 2001
From: Paul Mars <paul.mars@canonical.com>
Date: Tue, 26 Sep 2023 09:07:43 +0200
Subject: [PATCH] Cover new cases for customizeFstab and
 populateClassicRootfsContents

---
 internal/statemachine/classic_test.go | 63 ++++++++++++++++++++++++---
 1 file changed, 58 insertions(+), 5 deletions(-)

diff --git a/internal/statemachine/classic_test.go b/internal/statemachine/classic_test.go
index 3df09463..501b094c 100644
--- a/internal/statemachine/classic_test.go
+++ b/internal/statemachine/classic_test.go
@@ -1579,9 +1579,9 @@ func TestFailedPrepareClassicImage(t *testing.T) {
 	})
 }
 
-// TestPopulateClassicRootfsContents runs the state machine through populate_rootfs_contents and examines
+// TestStateMachine_PopulateClassicRootfsContents runs the state machine through populate_rootfs_contents and examines
 // the rootfs to ensure at least some of the correct file are in place
-func TestPopulateClassicRootfsContents(t *testing.T) {
+func TestStateMachine_PopulateClassicRootfsContents(t *testing.T) {
 	t.Run("test_populate_classic_rootfs_contents", func(t *testing.T) {
 		if runtime.GOARCH != "amd64" {
 			t.Skip("Test for amd64 only")
@@ -1599,6 +1599,7 @@ func TestPopulateClassicRootfsContents(t *testing.T) {
 			Rootfs: &imagedefinition.Rootfs{
 				Archive: "ubuntu",
 			},
+			Customization: &imagedefinition.Customization{},
 		}
 
 		// need workdir set up for this
@@ -1625,13 +1626,28 @@ func TestPopulateClassicRootfsContents(t *testing.T) {
 			}
 		}
 
+		// return when Customization.Fstab is not empty
+		stateMachine.ImageDef.Customization.Fstab = []*imagedefinition.Fstab{
+			{
+				Label:        "writable",
+				Mountpoint:   "/",
+				FSType:       "ext4",
+				MountOptions: "defaults",
+				Dump:         true,
+				FsckOrder:    1,
+			},
+		}
+
+		err = stateMachine.populateClassicRootfsContents()
+		asserter.AssertErrNil(err, true)
+
 		os.RemoveAll(stateMachine.stateMachineFlags.WorkDir)
 	})
 }
 
-// TestFailedPopulateClassicRootfsContents tests failed scenarios in populateClassicRootfsContents
+// TestStateMachine_FailedPopulateClassicRootfsContents tests failed scenarios in populateClassicRootfsContents
 // this is accomplished by mocking functions
-func TestFailedPopulateClassicRootfsContents(t *testing.T) {
+func TestStateMachine_FailedPopulateClassicRootfsContents(t *testing.T) {
 	t.Run("test_failed_populate_classic_rootfs_contents", func(t *testing.T) {
 		asserter := helper.Asserter{T: t}
 		var stateMachine ClassicStateMachine
@@ -1681,6 +1697,24 @@ func TestFailedPopulateClassicRootfsContents(t *testing.T) {
 		asserter.AssertErrContains(err, "Error writing to fstab")
 		osWriteFile = os.WriteFile
 
+		// mock os.ReadFile
+		osReadFile = mockReadFile
+		defer func() {
+			osReadFile = os.ReadFile
+		}()
+		err = stateMachine.populateClassicRootfsContents()
+		asserter.AssertErrContains(err, "Error reading fstab")
+		osReadFile = os.ReadFile
+
+		// return when existing fstab contains LABEL=writable
+		//nolint:gosec,G306
+		err = os.WriteFile(filepath.Join(stateMachine.tempDirs.chroot, "etc", "fstab"),
+			[]byte("LABEL=writable\n"),
+			0644)
+		asserter.AssertErrNil(err, true)
+		err = stateMachine.populateClassicRootfsContents()
+		asserter.AssertErrNil(err, true)
+
 		// create an /etc/resolv.conf.tmp in the chroot
 		err = os.MkdirAll(filepath.Join(stateMachine.tempDirs.chroot, "etc"), 0755)
 		asserter.AssertErrNil(err, true)
@@ -2963,6 +2997,7 @@ func TestCustomizeFstab(t *testing.T) {
 	testCases := []struct {
 		name          string
 		fstab         []*imagedefinition.Fstab
+		fstabTruncate bool
 		expectedFstab string
 		existingFstab string
 	}{
@@ -2995,6 +3030,23 @@ func TestCustomizeFstab(t *testing.T) {
 			},
 			expectedFstab: `LABEL=xxx / ext4 discard,errors=remount-ro 0 1
 LABEL=writable	/	ext4	defaults	1	1
+`,
+			existingFstab: `LABEL=xxx / ext4 discard,errors=remount-ro 0 1`,
+		},
+		{
+			name: "one_entry to a non-empty fstab to be truncated",
+			fstab: []*imagedefinition.Fstab{
+				{
+					Label:        "writable",
+					Mountpoint:   "/",
+					FSType:       "ext4",
+					MountOptions: "defaults",
+					Dump:         true,
+					FsckOrder:    1,
+				},
+			},
+			fstabTruncate: true,
+			expectedFstab: `LABEL=writable	/	ext4	defaults	1	1
 `,
 			existingFstab: `LABEL=xxx / ext4 discard,errors=remount-ro 0 1`,
 		},
@@ -3051,7 +3103,8 @@ LABEL=system-boot	/boot/firmware	vfat	defaults	0	1
 				Series:       getHostSuite(),
 				Rootfs:       &imagedefinition.Rootfs{},
 				Customization: &imagedefinition.Customization{
-					Fstab: tc.fstab,
+					Fstab:         tc.fstab,
+					FstabTruncate: tc.fstabTruncate,
 				},
 			}