More unit test fixes

Author: jscalev

internal/cleanup/cleanup.go

@@ -2,8 +2,11 @@ package cleanup
 
 import (
 	"fmt"
+	"os"
 	"path/filepath"
+	"strconv"
 
+	"github.com/Azure/azure-extension-platform/pkg/utils"
 	"github.com/Azure/run-command-handler-linux/internal/constants"
 	"github.com/Azure/run-command-handler-linux/internal/types"
 	"github.com/Azure/run-command-handler-linux/pkg/linuxutils"
@@ -38,28 +41,28 @@ func deleteAllScriptsAndSettings(ctx *log.Context, metadata types.RCMetadata, h
 }
 
 func deleteScriptsAndSettingsExceptMostRecent(ctx *log.Context, metadata types.RCMetadata, h types.HandlerEnvironment, runAsUser string) {
-	// runtimeSettingsRegexFormat := metadata.ExtName + ".\\d+.settings"
-	// runtimeSettingsLastSeqNumFormat := metadata.ExtName + ".%d.settings"
+	runtimeSettingsRegexFormat := metadata.ExtName + ".\\d+.settings"
+	runtimeSettingsLastSeqNumFormat := metadata.ExtName + ".%d.settings"
 
-	// // check if directory exists
-	// _, err := os.Open(metadata.DownloadPath)
-	// if err == nil {
-	// 	err := utils.TryClearExtensionScriptsDirectoriesAndSettingsFilesExceptMostRecent(metadata.DownloadPath, h.HandlerEnvironment.ConfigFolder, "",
-	// 		uint64(metadata.SeqNum), runtimeSettingsRegexFormat, runtimeSettingsLastSeqNumFormat)
-	// 	if err != nil {
-	// 		ctx.Log("event", "could not clear settings and script files", "error", err)
-	// 	}
-	// } else {
-	// 	ctx.Log("message", "directory does not exist. Skipping cleanup")
-	// }
+	// check if directory exists
+	_, err := os.Open(metadata.DownloadPath)
+	if err == nil {
+		err := utils.TryClearExtensionScriptsDirectoriesAndSettingsFilesExceptMostRecent(metadata.DownloadPath, h.HandlerEnvironment.ConfigFolder, "",
+			uint64(metadata.SeqNum), runtimeSettingsRegexFormat, runtimeSettingsLastSeqNumFormat)
+		if err != nil {
+			ctx.Log("event", "could not clear settings and script files", "error", err)
+		}
+	} else {
+		ctx.Log("message", "directory does not exist. Skipping cleanup")
+	}
 
-	// if runAsUser != "" {
-	// 	runAsDownloadParent := filepath.Join(fmt.Sprintf(constants.RunAsDir, runAsUser), metadata.DownloadDir)
-	// 	seqNumString := strconv.Itoa(metadata.SeqNum)
-	// 	ctx.Log("message", "removing all files from the download 'runas' directory "+runAsDownloadParent)
-	// 	err = utils.TryDeleteDirectoriesExcept(runAsDownloadParent, seqNumString)
-	// 	if err != nil {
-	// 		ctx.Log("event", "could not clear runas script")
-	// 	}
-	// }
+	if runAsUser != "" {
+		runAsDownloadParent := filepath.Join(fmt.Sprintf(constants.RunAsDir, runAsUser), metadata.DownloadDir)
+		seqNumString := strconv.Itoa(metadata.SeqNum)
+		ctx.Log("message", "removing all files from the download 'runas' directory "+runAsDownloadParent)
+		err = utils.TryDeleteDirectoriesExcept(runAsDownloadParent, seqNumString)
+		if err != nil {
+			ctx.Log("event", "could not clear runas script")
+		}
+	}
 }

internal/cmds/cmds.go

@@ -202,22 +202,22 @@ func enable(ctx *log.Context, h types.HandlerEnvironment, report *types.RunComma
 	}
 
 	dir := filepath.Join(metadata.DownloadPath, fmt.Sprintf("%d", metadata.SeqNum))
-	scriptFilePath, err := downloadScript(ctx, dir, &cfg)
-	if err != nil {
+	scriptFilePath, ewc := downloadScript(ctx, dir, &cfg)
+	if ewc != nil {
 		errMessage := fmt.Sprintf("Failed to download script: %v due to: %v", download.GetUriForLogging(cfg.ScriptURI()), err)
 		extensionEvents.LogErrorEvent("enable", errMessage)
 		return "",
 			"",
-			vmextension.CreateWrappedErrorWithClarification(err, fmt.Sprintf("File downloads failed. Use either a public script URI that points to .sh file, Azure storage blob SAS URI or storage blob accessible by a managed identity and retry. If managed identity is used, make sure it has been given access to container of storage blob '%s' with 'Storage Blob Data Reader' role assignment. In case of user-assigned identity, make sure you add it under VM's identity. For more info, refer https://aka.ms/RunCommandManagedLinux", download.GetUriForLogging(cfg.ScriptURI()))),
+			vmextension.CreateWrappedErrorWithClarification(ewc, fmt.Sprintf("File downloads failed. Use either a public script URI that points to .sh file, Azure storage blob SAS URI or storage blob accessible by a managed identity and retry. If managed identity is used, make sure it has been given access to container of storage blob '%s' with 'Storage Blob Data Reader' role assignment. In case of user-assigned identity, make sure you add it under VM's identity. For more info, refer https://aka.ms/RunCommandManagedLinux", download.GetUriForLogging(cfg.ScriptURI()))),
 			constants.FileDownload_GenericError
 	}
 
-	err = downloadArtifacts(ctx, dir, &cfg)
-	if err != nil {
-		errMessage := fmt.Sprintf("Failed to download artifacts: %v", err)
+	ewc = downloadArtifacts(ctx, dir, &cfg)
+	if ewc != nil {
+		errMessage := fmt.Sprintf("Failed to download artifacts: %v", ewc)
 		extensionEvents.LogErrorEvent("enable", errMessage)
 		return "", "",
-			vmextension.CreateWrappedErrorWithClarification(err, "Artifact downloads failed. Use either a public artifact URI that points to .sh file, Azure storage blob SAS URI, or storage blob accessible by a managed identity and retry."),
+			vmextension.CreateWrappedErrorWithClarification(ewc, "Artifact downloads failed. Use either a public artifact URI that points to .sh file, Azure storage blob SAS URI, or storage blob accessible by a managed identity and retry."),
 			constants.ArtifactDownload_GenericError
 	}
 

internal/cmds/cmds_test.go

@@ -12,11 +12,14 @@ import (
 	"strings"
 	"testing"
 
+	osexec "os/exec"
+
 	"github.com/Azure/azure-extension-platform/pkg/extensionevents"
 	"github.com/Azure/azure-extension-platform/pkg/handlerenv"
 	"github.com/Azure/azure-extension-platform/pkg/logging"
 	"github.com/Azure/azure-extension-platform/vmextension"
 	"github.com/Azure/run-command-handler-linux/internal/constants"
+	"github.com/Azure/run-command-handler-linux/internal/exec"
 	"github.com/Azure/run-command-handler-linux/internal/files"
 	"github.com/Azure/run-command-handler-linux/internal/handlersettings"
 	"github.com/Azure/run-command-handler-linux/internal/settings"
@@ -436,6 +439,13 @@ func Test_runCmd_success(t *testing.T) {
 	require.Nil(t, err)
 	defer os.RemoveAll(dir)
 
+	orig := exec.FnRunCommand
+	defer func() { exec.FnRunCommand = orig }()
+	exec.FnRunCommand = func(_ *osexec.Cmd) error {
+		// Success
+		return nil
+	}
+
 	metadata := types.NewRCMetadata("extName", 0, constants.DownloadFolder, DataDir)
 	err, exitCode := runCmd(log.NewContext(log.NewNopLogger()), dir, "", &handlersettings.HandlerSettings{
 		PublicSettings: handlersettings.PublicSettings{Source: &handlersettings.ScriptSource{Script: script}},

internal/commandProcessor/commandProcessor.go

@@ -101,7 +101,7 @@ func ProcessHandlerCommandWithDetails(ctx *log.Context, cmd types.Cmd, hEnv type
 		statusToReport := types.StatusSuccess
 
 		// Add an error clarification if we have one
-		var ewc vmextension.ErrorWithClarification
+		var ewc *vmextension.ErrorWithClarification
 		if errors.As(cmdInvokeError, &ewc) {
 			instView.ErrorClarificationValue = ewc.ErrorCode
 		}
@@ -113,6 +113,11 @@ func ProcessHandlerCommandWithDetails(ctx *log.Context, cmd types.Cmd, hEnv type
 		}
 
 		fnReportInstanceView(ctx, hEnv, metadata, statusToReport, cmd, &instView)
+
+		if err == nil {
+			return nil
+		}
+
 		return errors.Wrapf(err, "command execution failed")
 	} else { // No error. Succeeded
 		instView.ExecutionMessage = "Execution completed"

internal/commandProcessor/commandProcessor_test.go

@@ -202,20 +202,15 @@ func Test_ProcessHandlerCommandWithDetails_Failure_WithClarification(t *testing.
 		return nil
 	}
 
-	ewc := vmextension.ErrorWithClarification{
-		ErrorCode: 1234,
-		Err:       errors.New("the chipmunks are upset"),
-	}
-
 	mockFunc := types.CmdFunctions{
 		Invoke: func(_ *log.Context, _ types.HandlerEnvironment, iv *types.RunCommandInstanceView, _ types.RCMetadata, _ types.Cmd) (string, string, error, int) {
-			return "x", "y", ewc, 3
+			return "x", "y", vmextension.NewErrorWithClarificationPtr(1234, errors.New("the chipmunks are upset")), 3
 		},
 	}
 
 	orig := fnGetHandlerSettings
 	defer func() { fnGetHandlerSettings = orig }()
-	fnGetHandlerSettings = func(string, string, int, *log.Context) (handlersettings.HandlerSettings, error) {
+	fnGetHandlerSettings = func(string, string, int, *log.Context) (handlersettings.HandlerSettings, *vmextension.ErrorWithClarification) {
 		return handlersettings.HandlerSettings{
 			PublicSettings: handlersettings.PublicSettings{
 				TreatFailureAsDeploymentFailure: false,

internal/exec/exec.go

@@ -28,7 +28,7 @@ var (
 	fnOsMkDirAll = os.MkdirAll
 	fnOsOpenFile = os.OpenFile
 	fnOsSetEnv   = os.Setenv
-	fnRunCommand = runCommand
+	FnRunCommand = runCommand
 	fnUserLookup = user.Lookup
 )
 
@@ -140,7 +140,7 @@ func Exec(ctx *log.Context, cmd, workdir string, stdout, stderr io.WriteCloser,
 	command.Dir = workdir
 	command.Stdout = stdout
 	command.Stderr = stderr
-	err = fnRunCommand(command)
+	err = FnRunCommand(command)
 	if err != nil {
 		exitErr, ok := err.(*exec.ExitError)
 		if ok {

internal/exec/exec_test.go

@@ -47,7 +47,7 @@ func TestExec_AlwaysClosesStreams(t *testing.T) {
 	errw := newCloseRecorder()
 
 	// Force runCommand to return error to ensure closures happen on error path.
-	fnRunCommand = func(_ *exec.Cmd) error {
+	FnRunCommand = func(_ *exec.Cmd) error {
 		return errors.New("the chipmunks have revolted")
 	}
 
@@ -85,7 +85,7 @@ func TestExec_RunAsUser_OpenSourceScriptFails_ReturnsOpenSourceFailed(t *testing
 		return nil, errors.New("open failed")
 	}
 	// Ensure we don't accidentally reach execution
-	fnRunCommand = func(_ *exec.Cmd) error {
+	FnRunCommand = func(_ *exec.Cmd) error {
 		t.Fatalf("fnRunCommand should not be called when RunAs setup fails")
 		return nil
 	}
@@ -114,7 +114,7 @@ func TestExec_RunAsUser_CreateDestScriptFails_ReturnsOpenSourceFailed(t *testing
 	fnOsCreate = func(_ string) (*os.File, error) {
 		return nil, errors.New("create failed")
 	}
-	fnRunCommand = func(_ *exec.Cmd) error {
+	FnRunCommand = func(_ *exec.Cmd) error {
 		t.Fatalf("fnRunCommand should not be called when RunAs setup fails")
 		return nil
 	}
@@ -147,7 +147,7 @@ func TestExec_RunAsUser_CopyFails_ReturnsCopyFailed(t *testing.T) {
 	fnIoCopy = func(_ io.Writer, _ io.Reader) (int64, error) {
 		return 0, errors.New("the chipmunks do not copy")
 	}
-	fnRunCommand = func(_ *exec.Cmd) error {
+	FnRunCommand = func(_ *exec.Cmd) error {
 		t.Fatalf("fnRunCommand should not be called when RunAs setup fails")
 		return nil
 	}
@@ -179,7 +179,7 @@ func TestExec_RunAsUser_LookupUserFails_ReturnsRunAsUserLogonFailed(t *testing.T
 	fnUserLookup = func(_ string) (*user.User, error) {
 		return nil, errors.New("no such chipmunk")
 	}
-	fnRunCommand = func(_ *exec.Cmd) error {
+	FnRunCommand = func(_ *exec.Cmd) error {
 		t.Fatalf("fnRunCommand should not be called when RunAs setup fails")
 		return nil
 	}
@@ -210,7 +210,7 @@ func TestExec_RunAsUser_UidParseFails_ReturnsLookupUserUidFailed(t *testing.T) {
 	fnUserLookup = func(_ string) (*user.User, error) {
 		return &user.User{Uid: "not-an-int"}, nil
 	}
-	fnRunCommand = func(_ *exec.Cmd) error {
+	FnRunCommand = func(_ *exec.Cmd) error {
 		t.Fatalf("fnRunCommand should not be called when RunAs setup fails")
 		return nil
 	}
@@ -243,7 +243,7 @@ func TestExec_RunAsUser_ChownFails_ReturnsChangeOwnerFailed(t *testing.T) {
 	fnOsChown = func(_ string, _ int, _ int) error {
 		return errors.New("chown failed")
 	}
-	fnRunCommand = func(_ *exec.Cmd) error {
+	FnRunCommand = func(_ *exec.Cmd) error {
 		t.Fatalf("fnRunCommand should not be called when RunAs setup fails")
 		return nil
 	}
@@ -278,7 +278,7 @@ func TestExec_RunAsUser_ChmodFails_ReturnsChangePermissionsFailed(t *testing.T)
 		return errors.New("chmod failed")
 	}
 
-	fnRunCommand = func(_ *exec.Cmd) error {
+	FnRunCommand = func(_ *exec.Cmd) error {
 		t.Fatalf("fnRunCommand should not be called when RunAs setup fails")
 		return nil
 	}
@@ -539,7 +539,7 @@ func saveAndRestoreFns() func() {
 		mkdirAll:   fnOsMkDirAll,
 		openFile:   fnOsOpenFile,
 		setEnv:     fnOsSetEnv,
-		runCommand: fnRunCommand,
+		runCommand: FnRunCommand,
 		userLookup: fnUserLookup,
 	}
 	return func() {
@@ -550,7 +550,7 @@ func saveAndRestoreFns() func() {
 		fnOsMkDirAll = s.mkdirAll
 		fnOsOpenFile = s.openFile
 		fnOsSetEnv = s.setEnv
-		fnRunCommand = s.runCommand
+		FnRunCommand = s.runCommand
 		fnUserLookup = s.userLookup
 	}
 }

internal/files/files.go

@@ -55,11 +55,11 @@ func downloadAndProcessURL(ctx *log.Context, url, downloadDir string, fileName s
 
 	targetFilePath := filepath.Join(downloadDir, fileName)
 
-	var scriptSASDownloadErr error = nil
+	var scriptSASDownloadErr *vmextension.ErrorWithClarification = nil
 	var downloadedFilePath string = ""
 	if scriptSAS != "" {
 		if UseMockSASDownloadFailure {
-			scriptSASDownloadErr = errors.New("Downloading script using SAS token failed.")
+			scriptSASDownloadErr = vmextension.NewErrorWithClarificationPtr(42, errors.New("Downloading script using SAS token failed."))
 		} else {
 			downloadedFilePath, scriptSASDownloadErr = download.GetSASBlob(url, scriptSAS, downloadDir)
 		}

internal/handlersettings/handlersettings.go

@@ -1,6 +1,7 @@
 package handlersettings
 
 import (
+	"github.com/Azure/azure-extension-platform/vmextension"
 	"github.com/go-kit/kit/log"
 	"github.com/pkg/errors"
 )
@@ -11,7 +12,7 @@ var (
 
 // parseAndValidateSettings reads configuration from configFolder, decrypts it,
 // runs JSON-schema and logical validation on it and returns it back.
-func ParseAndValidateSettings(ctx *log.Context, configFilePath string) (h HandlerSettings, _ error) {
+func ParseAndValidateSettings(ctx *log.Context, configFilePath string) (h HandlerSettings, _ *vmextension.ErrorWithClarification) {
 	ctx.Log("event", "reading configuration from "+configFilePath)
 	pubJSON, protJSON, err := readSettings(configFilePath)
 	if err != nil {
@@ -36,7 +37,7 @@ func ParseAndValidateSettings(ctx *log.Context, configFilePath string) (h Handle
 // readSettings uses specified configFilePath (comes from HandlerEnvironment) to
 // decrypt and parse the public/protected settings of the extension handler into
 // JSON objects.
-func readSettings(configFilePath string) (pubSettingsJSON, protSettingsJSON map[string]interface{}, err error) {
+func readSettings(configFilePath string) (pubSettingsJSON, protSettingsJSON map[string]interface{}, err *vmextension.ErrorWithClarification) {
 	pubSettingsJSON, protSettingsJSON, err = ReadSettings(configFilePath)
 	return
 }

internal/handlersettings/utilities.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 
+	"github.com/Azure/azure-extension-platform/vmextension"
 	"github.com/go-kit/kit/log"
 )
 
@@ -35,7 +36,7 @@ func GetUriForLogging(uriString string) string {
 }
 
 // Get handler settings from config folder. Example path: /var/lib/waagent/Microsoft.CPlat.Core.RunCommandHandlerLinux-1.3.2/config
-func GetHandlerSettings(configFolder string, extensionName string, sequenceNumber int, logContext *log.Context) (HandlerSettings, error) {
+func GetHandlerSettings(configFolder string, extensionName string, sequenceNumber int, logContext *log.Context) (HandlerSettings, *vmextension.ErrorWithClarification) {
 	configPath := GetConfigFilePath(configFolder, sequenceNumber, extensionName)
 	cfg, err := ParseAndValidateSettings(logContext, configPath)
 	return cfg, err