WIP

internal/imagedefinition/README.rst

@@ -251,6 +251,14 @@ The following specification defines what is supported in the YAML:
              -
                # Path inside the rootfs.
                path: <string>
+               # Arguments to give to the command
+               args: (optional)
+                 - <string>
+                 - <string>
+               # Environment variables to set before executing the command
+               env: (optional)
+                 - <string>
+                 - <string>
            # Any additional users to add in the rootfs
            # We recommend using cloud-init when possible and fallback
            # on this method if not possible (e.g performance issues)

internal/imagedefinition/image_definition.go

@@ -145,9 +145,11 @@ type CopyFile struct {
 	Source string `yaml:"source"      json:"Source"`
 }
 
-// Execute allows users to execute a script in the rootfs of an image
+// Execute allows users to execute a script/command in the rootfs of an image
 type Execute struct {
-	ExecutePath string `yaml:"path" json:"ExecutePath"`
+	ExecutePath string   `yaml:"path" json:"ExecutePath"`
+	ExecuteArgs []string `yaml:"args" json:"ExecuteArgs,omitempty"`
+	Env         []string `yaml:"env"  json:"Env,omitempty"`
 }
 
 // TouchFile allows users to touch a file in the rootfs of an image

internal/statemachine/classic_test.go

@@ -3818,7 +3818,7 @@ func TestStateMachine_installPackages_checkcmds(t *testing.T) {
 
 	t.Cleanup(func() { os.RemoveAll(stateMachine.stateMachineFlags.WorkDir) })
 
-	mockCmder := NewMockExecCommand()
+	mockCmder := NewMockExecCommander()
 
 	execCommand = mockCmder.Command
 	t.Cleanup(func() { execCommand = exec.Command })
@@ -3886,7 +3886,7 @@ func TestStateMachine_installPackages_checkcmds_failing(t *testing.T) {
 
 	t.Cleanup(func() { os.RemoveAll(stateMachine.stateMachineFlags.WorkDir) })
 
-	mockCmder := NewMockExecCommand()
+	mockCmder := NewMockExecCommander()
 
 	execCommand = mockCmder.Command
 	t.Cleanup(func() { execCommand = exec.Command })

internal/statemachine/helper.go

@@ -870,10 +870,12 @@ func manualCopyFile(customizations []*imagedefinition.CopyFile, confDefPath stri
 // manualExecute executes executable files in the chroot
 func manualExecute(customizations []*imagedefinition.Execute, targetDir string, debug bool) error {
 	for _, c := range customizations {
-		executeCmd := execCommand("chroot", targetDir, c.ExecutePath)
+		executeCmd := execCommand("chroot", append([]string{targetDir, c.ExecutePath}, c.ExecuteArgs...)...)
 		if debug {
 			fmt.Printf("Executing command \"%s\"\n", executeCmd.String())
 		}
+		executeCmd.Env = append(executeCmd.Env, c.Env...)
+
 		executeOutput := helper.SetCommandOutput(executeCmd, debug)
 		err := executeCmd.Run()
 		if err != nil {

internal/statemachine/helper_test.go

@@ -782,6 +782,125 @@ func TestFailedManualTouchFile(t *testing.T) {
 	asserter.AssertErrContains(err, "Error creating file")
 }
 
+// TestStateMachine_manualExecute tests manualExecute
+func TestStateMachine_manualExecute(t *testing.T) {
+	type cmdMatcher struct {
+		cmdRegex *regexp.Regexp
+		env      []string
+	}
+
+	type args struct {
+		customizations []*imagedefinition.Execute
+		targetDir      string
+		debug          bool
+	}
+	testCases := []struct {
+		name        string
+		args        args
+		expected    []cmdMatcher
+		expectedErr string
+	}{
+		{
+			name: "single simple command",
+			args: args{
+				customizations: []*imagedefinition.Execute{
+					{
+						ExecutePath: "/execute/path",
+						ExecuteArgs: []string{"arg1", "arg2"},
+						Env:         []string{"VAR1=value1", "VAR2=value2"},
+					},
+				},
+				targetDir: "test",
+				debug:     true,
+			},
+			expected: []cmdMatcher{
+				{
+					cmdRegex: regexp.MustCompile("Executing command .*"),
+				},
+				{
+					cmdRegex: regexp.MustCompile("chroot test /execute/path arg1 arg2"),
+					env: []string{
+						"",
+					},
+				},
+			},
+		},
+		{
+			name: "3 commands",
+			args: args{
+				customizations: []*imagedefinition.Execute{
+					{
+						ExecutePath: "/execute/path1",
+						ExecuteArgs: []string{"arg1", "arg2"},
+						Env:         []string{"VAR1=value1", "VAR2=value2"},
+					},
+					{
+						ExecutePath: "/execute/path2",
+						ExecuteArgs: []string{"arg21", "arg22"},
+						Env:         []string{"VAR1=value21", "VAR2=value22"},
+					},
+					{
+						ExecutePath: "/execute/path3",
+						ExecuteArgs: []string{"arg31", "arg32"},
+						Env:         []string{"VAR1=value31", "VAR2=value32"},
+					},
+				},
+				targetDir: "test",
+				debug:     true,
+			},
+			expected: []cmdMatcher{
+				{
+					cmdRegex: regexp.MustCompile("Executing command .*"),
+				},
+				{
+					cmdRegex: regexp.MustCompile("chroot test /execute/path arg1 arg2"),
+					env: []string{
+						"",
+					},
+				},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+
+			asserter := helper.Asserter{T: t}
+
+			mockCmder := NewMockExecCommander()
+
+			execCommand = mockCmder.Command
+			t.Cleanup(func() { execCommand = exec.Command })
+
+			stdout, restoreStdout, _ := helper.CaptureStd(&os.Stdout)
+			t.Cleanup(func() { restoreStdout() })
+
+			err := manualExecute(tc.args.customizations, tc.args.targetDir, tc.args.debug)
+			asserter.AssertErrNil(err, true)
+
+			restoreStdout()
+			readStdout, _ := io.ReadAll(stdout)
+
+			gotCmds := strings.Split(strings.TrimSpace(string(readStdout)), "\n")
+			if len(tc.expected) != len(gotCmds) {
+				t.Fatalf("%v commands to be executed, expected %v", len(gotCmds), len(tc.expected))
+			}
+
+			for i, gotCmd := range gotCmds {
+				expectedCmd := tc.expected[i].cmdRegex
+				if !expectedCmd.Match([]byte(gotCmd)) {
+					t.Errorf("Cmd \"%v\" not matching. Expected %v\n", gotCmd, expectedCmd.String())
+				}
+			}
+
+			for i, gotCmd := range mockCmder.cmds {
+				expectedEnv := tc.args.customizations[i].Env
+				asserter.AssertEqual(expectedEnv, gotCmd.Env)
+			}
+		})
+	}
+}
+
 // TestFailedManualExecute tests the fail case of the manualExecute function
 func TestFailedManualExecute(t *testing.T) {
 	t.Parallel()
@@ -1366,7 +1485,7 @@ func TestStateMachine_updateGrub_checkcmds(t *testing.T) {
 
 	t.Cleanup(func() { os.RemoveAll(stateMachine.stateMachineFlags.WorkDir) })
 
-	mockCmder := NewMockExecCommand()
+	mockCmder := NewMockExecCommander()
 
 	execCommand = mockCmder.Command
 	t.Cleanup(func() { execCommand = exec.Command })

internal/statemachine/tests_helper_test.go

@@ -192,14 +192,32 @@ func (m *mockRunCmd) runCmd(cmd *exec.Cmd, debug bool) error {
 	return nil
 }
 
-type mockExecCmd struct{}
-
 func NewMockExecCommand() *mockExecCmd {
 	return &mockExecCmd{}
 }
 
-func (m *mockExecCmd) Command(cmd string, args ...string) *exec.Cmd {
+type mockExecCmder struct {
+	cmds []*exec.Cmd
+}
+
+type mockExecCmd struct {
+	exec.Cmd
+	called bool
+}
+
+func (m *mockExecCmd) Run() error {
+	m.called = true
+	return m.Run()
+}
+
+func NewMockExecCommander() *mockExecCmder {
+	return &mockExecCmder{}
+}
+
+func (m *mockExecCmder) Command(name string, args ...string) *exec.Cmd {
 	// Replace the command with an echo of it
 	//nolint:gosec,G204
-	return exec.Command("echo", append([]string{cmd}, args...)...)
+	cmd := exec.Command("echo", append([]string{name}, args...)...)
+	m.cmds = append(m.cmds, cmd)
+	return cmd
 }