From 71bba8ebcfcc0e7c6b254711ea064aadd965edee Mon Sep 17 00:00:00 2001 From: Paul Mars Date: Mon, 22 Jan 2024 09:38:22 +0100 Subject: [PATCH] WIP --- internal/imagedefinition/README.rst | 8 ++ internal/imagedefinition/image_definition.go | 6 +- internal/statemachine/classic_test.go | 4 +- internal/statemachine/helper.go | 4 +- internal/statemachine/helper_test.go | 121 ++++++++++++++++++- internal/statemachine/tests_helper_test.go | 26 +++- 6 files changed, 159 insertions(+), 10 deletions(-) diff --git a/internal/imagedefinition/README.rst b/internal/imagedefinition/README.rst index 8a7d2d88..dc204fe1 100644 --- a/internal/imagedefinition/README.rst +++ b/internal/imagedefinition/README.rst @@ -251,6 +251,14 @@ The following specification defines what is supported in the YAML: - # Path inside the rootfs. path: + # Arguments to give to the command + args: (optional) + - + - + # Environment variables to set before executing the command + env: (optional) + - + - # 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) diff --git a/internal/imagedefinition/image_definition.go b/internal/imagedefinition/image_definition.go index 7ea65659..3bedc334 100644 --- a/internal/imagedefinition/image_definition.go +++ b/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 diff --git a/internal/statemachine/classic_test.go b/internal/statemachine/classic_test.go index 21609072..35de3a52 100644 --- a/internal/statemachine/classic_test.go +++ b/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 }) diff --git a/internal/statemachine/helper.go b/internal/statemachine/helper.go index 40832181..5179ab57 100644 --- a/internal/statemachine/helper.go +++ b/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 { diff --git a/internal/statemachine/helper_test.go b/internal/statemachine/helper_test.go index 27312ebb..08499a1c 100644 --- a/internal/statemachine/helper_test.go +++ b/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 }) diff --git a/internal/statemachine/tests_helper_test.go b/internal/statemachine/tests_helper_test.go index 9bd53c49..19081049 100644 --- a/internal/statemachine/tests_helper_test.go +++ b/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 }