Added function to retrieve Environment variables set by the guest agent.

pkg/lockedfile/lockedfile_test.go

@@ -31,6 +31,7 @@ func TestLockFileMetadata(t *testing.T) {
 	var lf ILockedFile
 	lf, err := New(testFilePath, time.Second)
 	assert.NoError(t, err)
+	time.Sleep(time.Second)
 	lf.Close()
 	lastOpened, lastClosed, err := getLastOpenedAndLastClosedTime(testFilePath)
 	assert.NoError(t, err)

vmextension/vmextension.go

@@ -23,6 +23,22 @@ import (
 // Azure Guest Agent.
 const handlerEnvFileName = "HandlerEnvironment.json"
 
+// GuestAgentEnvVar represents environment variable names set by the Azure Guest Agent.
+type GuestAgentEnvVar string
+
+const (
+	// environment variables are declared here
+	// https://github.com/Azure/azure-vmextension-publishing/wiki/2.0-Partner-Guide-Handler-Design-Details#236-summary
+	GuestAgentEnvVarExtensionVersion     GuestAgentEnvVar = "AZURE_GUEST_AGENT_EXTENSION_VERSION"
+	GuestAgentEnvVarExtensionFullPath    GuestAgentEnvVar = "AZURE_GUEST_AGENT_EXTENSION_PATH"
+	GuestAgentEnvVarConfigSequenceNumber GuestAgentEnvVar = "ConfigSequenceNumber"
+	GuestAgentEnvVarUpdateToVersion      GuestAgentEnvVar = "VERSION"
+	GuestAgentEnvVarUpdateFromVersion    GuestAgentEnvVar = "AZURE_GUEST_AGENT_UPDATING_FROM_VERSION"
+	GuestAgentEnvVarDisableCmdExitCode   GuestAgentEnvVar = "AZURE_GUEST_AGENT_DISABLE_CMD_EXIT_CODE"
+	GuestAgentEnvVarUninstallCmdExitCode GuestAgentEnvVar = "AZURE_GUEST_AGENT_UNINSTALL_CMD_EXIT_CODE"
+	// environment variables that are for multiconfig extensions are excluded
+)
+
 type OperationName string
 
 const (
@@ -128,6 +144,14 @@ func (*prodGetVMExtensionEnvironmentManager) SetSequenceNumberInternal(extension
 	return seqno.SetSequenceNumber(extensionName, extensionVersion, seqNo)
 }
 
+func GetGuestAgentEnvironmetVariable(envVarName GuestAgentEnvVar) (string, error) {
+	extensionVersion, isSet := os.LookupEnv(string(envVarName))
+	if !isSet || extensionVersion == "" {
+		return "", errors.Errorf("Extension environment variable '%s' is not set", envVarName)
+	}
+	return extensionVersion, nil
+}
+
 // GetVMExtension returns a new VMExtension object
 func GetVMExtension(initInfo *InitializationInfo) (ext *VMExtension, _ error) {
 	return getVMExtensionInternal(initInfo, &prodGetVMExtensionEnvironmentManager{})

vmextension/vmextension_test.go

@@ -723,3 +723,23 @@ func cleanupDirsForVMExtension(vmExt *VMExtension) (combinedError error) {
 	combinedError = extensionerrors.CombineErrors(combinedError, os.RemoveAll(vmExt.HandlerEnv.DataFolder))
 	return
 }
+
+func Test_GuestAgentEnvironmentVariable_ReturnsValue(t *testing.T) {
+	t.Setenv(string(GuestAgentEnvVarExtensionVersion), "1.2.3")
+	version, err := GetGuestAgentEnvironmetVariable(GuestAgentEnvVarExtensionVersion)
+	require.NoError(t, err)
+	require.Equal(t, "1.2.3", version)
+}
+
+func Test_GuestAgentEnvironmentVariable_EmptyValue(t *testing.T) {
+	t.Setenv(string(GuestAgentEnvVarExtensionVersion), "")
+	_, err := GetGuestAgentEnvironmetVariable(GuestAgentEnvVarExtensionVersion)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), string(GuestAgentEnvVarExtensionVersion))
+}
+
+func Test_GuestAgentEnvironmentVariable_Unset(t *testing.T) {
+	os.Unsetenv(string(GuestAgentEnvVarExtensionVersion))
+	_, err := GetGuestAgentEnvironmetVariable(GuestAgentEnvVarExtensionVersion)
+	require.Error(t, err)
+}