Introduced CustomStatusFormatter to InitializationInfo (#17)

* Created a status formatter override function

* made variable private, added unit test

* made test more stable

* added assignment for customStatusFormatter

* Updated comment

* Addressed PR comments

* trigger new pipeline build

pkg/extensionevents/extension_events_linux.go

@@ -1,9 +1,11 @@
 package extensionevents
 
 import (
+	"fmt"
+
 	"golang.org/x/sys/unix"
 )
 
 func getThreadID() string {
-	return string(unix.Gettid())
+	return fmt.Sprintf("%d", unix.Gettid())
 }

pkg/seqno/seqno_test.go

@@ -64,11 +64,11 @@ func Test_findSeqNoFilesInDifferentOrder(t *testing.T) {
 	// findSeqNum should return the most recently created file, sleep it necessary to ensure that the creation times
 	// are different enough
 	writeSequenceNumberFile(t, sequenceNumberTestFolder, "5")
-	time.Sleep(5 * time.Millisecond)
+	time.Sleep(10 * time.Millisecond)
 	writeSequenceNumberFile(t, sequenceNumberTestFolder, "4")
-	time.Sleep(5 * time.Millisecond)
+	time.Sleep(10 * time.Millisecond)
 	writeSequenceNumberFile(t, sequenceNumberTestFolder, "3")
-	time.Sleep(5 * time.Millisecond)
+	time.Sleep(10 * time.Millisecond)
 	writeSequenceNumberFile(t, sequenceNumberTestFolder, "2")
 
 	seqNo, err := FindSeqNum(el, sequenceNumberTestFolder)

pkg/status/status.go

@@ -11,8 +11,8 @@ import (
 
 const (
 	// extension needn't write status files for other operations
-	EnableStatus = "Enable"
-	UpdateStatus = "Update"
+	EnableStatus  = "Enable"
+	UpdateStatus  = "Update"
 	DisableStatus = "Disable"
 )
 
@@ -110,13 +110,15 @@ func (r StatusReport) Save(statusFolder string, seqNo uint) error {
 	return nil
 }
 
+type StatusMessageFormatter func(operationName string, t StatusType, msg string) string
+
 // StatusMsg creates the reported status message based on the provided operation
 // type and the given message string.
 //
 // A message will be generated for empty string. For error status, pass the
 // error message.
-func StatusMsg(o string, t StatusType, msg string) string {
-	s := o
+func StatusMsg(operationName string, t StatusType, msg string) string {
+	s := operationName
 	switch t {
 	case StatusSuccess:
 		s += " succeeded"

vmextension/initialization.go

@@ -2,6 +2,7 @@ package vmextension
 
 import (
 	"github.com/Azure/azure-extension-platform/pkg/extensionerrors"
+	"github.com/Azure/azure-extension-platform/pkg/status"
 )
 
 // CallbackFunc is used for a non-Enable operation callback
@@ -12,19 +13,20 @@ type EnableCallbackFunc func(ext *VMExtension) (string, error)
 
 // InitializationInfo is passed by the extension to specify how the framework should run
 type InitializationInfo struct {
-	Name                string             // The name of the extension, without the Linux or Windows suffix
-	Version             string             // The version of the extension
-	SupportsDisable     bool               // True if we should automatically disable the extension if Disable is called
-	SupportsResetState  bool               // True if we should remove all contents of all folder when ResetState is called
-	RequiresSeqNoChange bool               // True if Enable will only execute if the sequence number changes
-	InstallExitCode     int                // Exit code to use for the install case
-	OtherExitCode       int                // Exit code to use for all other cases
-	EnableCallback      EnableCallbackFunc // Called for the enable operation
-	DisableCallback     CallbackFunc       // Called for the Disable operation. Only set this if the extension wants a callback.
-	UpdateCallback      CallbackFunc       // Called for the Update operation. If nil, then update is not supported.
-	ResetStateCallback  CallbackFunc       // Called for the ResetState operation. Only set this if the extension wants a callback.
-	InstallCallback     CallbackFunc       // Called for the Install operation. Only set this if the extension wants a callback.
-	UninstallCallback   CallbackFunc       // Called for the Uninstall operation. Only set this if the extension wants a callback.
+	Name                  string                        // The name of the extension, without the Linux or Windows suffix
+	Version               string                        // The version of the extension
+	SupportsDisable       bool                          // True if we should automatically disable the extension if Disable is called
+	SupportsResetState    bool                          // True if we should remove all contents of all folder when ResetState is called
+	RequiresSeqNoChange   bool                          // True if Enable will only execute if the sequence number changes
+	InstallExitCode       int                           // Exit code to use for the install case
+	OtherExitCode         int                           // Exit code to use for all other cases
+	EnableCallback        EnableCallbackFunc            // Called for the enable operation
+	DisableCallback       CallbackFunc                  // Called for the Disable operation. Only set this if the extension wants a callback.
+	UpdateCallback        CallbackFunc                  // Called for the Update operation. If nil, then update is not supported.
+	ResetStateCallback    CallbackFunc                  // Called for the ResetState operation. Only set this if the extension wants a callback.
+	InstallCallback       CallbackFunc                  // Called for the Install operation. Only set this if the extension wants a callback.
+	UninstallCallback     CallbackFunc                  // Called for the Uninstall operation. Only set this if the extension wants a callback.
+	CustomStatusFormatter status.StatusMessageFormatter // Provide a function to format the status message. If nil default formatting behavior will be preserved.
 }
 
 // GetInitializationInfo returns a new InitializationInfo object

vmextension/vmextension.go

@@ -85,25 +85,6 @@ type executionInfo struct {
 	manager             environmentmanager.IGetVMExtensionEnvironmentManager // Used by tests to mock the environment
 }
 
-// HandlerEnvironment describes the handler environment configuration presented
-// to the extension handler by the Azure Guest Agent.
-type handlerEnvironmentInternal struct {
-	Version            float64 `json:"version"`
-	Name               string  `json:"name"`
-	HandlerEnvironment struct {
-		HeartbeatFile       string `json:"heartbeatFile"`
-		StatusFolder        string `json:"statusFolder"`
-		ConfigFolder        string `json:"configFolder"`
-		LogFolder           string `json:"logFolder"`
-		EventsFolder        string `json:"eventsFolder"`
-		EventsFolderPreview string `json:"eventsFolder_preview"`
-		DeploymentID        string `json:"deploymentid"`
-		RoleName            string `json:"rolename"`
-		Instance            string `json:"instance"`
-		HostResolverAddress string `json:"hostResolverAddress"`
-	}
-}
-
 // VMExtension is an abstraction for standard extension operations in an OS agnostic manner
 type VMExtension struct {
 	Name                       string                                    // The name of the extension. This will contain 'Windows' or 'Linux'
@@ -115,6 +96,7 @@ type VMExtension struct {
 	ExtensionEvents            *extensionevents.ExtensionEventManager    // Allows extensions to raise events
 	ExtensionLogger            *logging.ExtensionLogger                  // Automatically logs to the log directory
 	exec                       *executionInfo                            // Internal information necessary for the extension to run
+	statusFormatter            status.StatusMessageFormatter             // Custom status message formatter from initialization info
 }
 
 type prodGetVMExtensionEnvironmentManager struct {
@@ -189,7 +171,7 @@ func getVMExtensionInternal(initInfo *InitializationInfo, manager environmentman
 			// current sequence number could not be found, this is a special error
 			currentSeqNo = nil
 		} else {
-			return nil, fmt.Errorf("Failed to read the current sequence number due to '%v'", err)
+			return nil, fmt.Errorf("failed to read the current sequence number due to '%v'", err)
 		}
 	} else {
 		*currentSeqNo = retrievedSequenceNumber
@@ -225,6 +207,13 @@ func getVMExtensionInternal(initInfo *InitializationInfo, manager environmentman
 		return manager.GetHandlerSettings(extensionLogger, handlerEnv)
 	}
 
+	var statusFormatter status.StatusMessageFormatter
+	if initInfo.CustomStatusFormatter != nil {
+		statusFormatter = initInfo.CustomStatusFormatter
+	} else {
+		statusFormatter = status.StatusMsg
+	}
+
 	ext = &VMExtension{
 		Name:                       initInfo.Name,
 		Version:                    initInfo.Version,
@@ -234,6 +223,7 @@ func getVMExtensionInternal(initInfo *InitializationInfo, manager environmentman
 		GetSettings:                settings,
 		ExtensionEvents:            extensionEvents,
 		ExtensionLogger:            extensionLogger,
+		statusFormatter:            statusFormatter,
 		exec: &executionInfo{
 			manager:             manager,
 			requiresSeqNoChange: initInfo.RequiresSeqNoChange,
@@ -288,7 +278,7 @@ func reportStatus(ve *VMExtension, t status.StatusType, c cmd, msg string) error
 		return err
 	}
 
-	s := status.New(t, c.operation.ToStatusName(), status.StatusMsg(c.operation.ToStatusName(), t, msg))
+	s := status.New(t, c.operation.ToStatusName(), ve.statusFormatter(c.operation.ToStatusName(), t, msg))
 	if err := s.Save(ve.HandlerEnv.StatusFolder, requestedSequenceNumber); err != nil {
 		ve.ExtensionLogger.Error("Failed to save handler status: %v", err)
 		return errors.Wrap(err, "failed to save handler status")
@@ -323,7 +313,7 @@ func (ve *VMExtension) printUsage(args []string) {
 	fmt.Printf("Usage: %s ", os.Args[0])
 	i := 0
 	for k := range ve.exec.cmds {
-		fmt.Printf(k.ToString())
+		fmt.Print(k.ToString())
 		if i != len(ve.exec.cmds)-1 {
 			fmt.Printf("|")
 		}

vmextension/vmextension_test.go

@@ -121,6 +121,30 @@ func Test_reportStatusSaved(t *testing.T) {
 	require.NoError(t, err, "File doesn't exist")
 }
 
+func Test_reportStatusFormatter(t *testing.T) {
+	ext := createTestVMExtension()
+
+	c := cmd{nil, InstallOperation, true, 99}
+	ext.HandlerEnv.StatusFolder = statusTestDirectory
+	ext.GetRequestedSequenceNumber = func() (uint, error) { return 45, nil }
+
+	createDirsForVMExtension(ext)
+	defer cleanupDirsForVMExtension(ext)
+
+	customFormattedMessage := "I am custom message"
+	ext.statusFormatter = func(operationName string, t status.StatusType, msg string) string {
+		return customFormattedMessage
+	}
+	err := reportStatus(ext, status.StatusSuccess, c, "msg")
+	require.NoError(t, err, "reportStatus failed")
+	statusFilePath := path.Join(statusTestDirectory, "45.status")
+	_, err = os.Stat(statusFilePath)
+	require.NoError(t, err, "File doesn't exist")
+	statusFileBytes, err := ioutil.ReadFile(statusFilePath)
+	require.NoError(t, err, "Could not read status file contents")
+	require.Contains(t, string(statusFileBytes), customFormattedMessage)
+}
+
 func Test_getVMExtensionNilValues(t *testing.T) {
 	_, err := GetVMExtension(nil)
 	require.Equal(t, extensionerrors.ErrArgCannotBeNull, err)
@@ -604,6 +628,7 @@ func createTestVMExtension() *VMExtension {
 			updateCallback:      nil,
 			cmds:                map[OperationName]cmd{DisableOperation: disableCommand},
 		},
+		statusFormatter: status.StatusMsg,
 	}
 }