From 777c87471b9d51b252c35f552bfa472f5f3e0417 Mon Sep 17 00:00:00 2001
From: r2k1 <yokree@gmail.com>
Date: Tue, 29 Oct 2024 19:53:48 +1300
Subject: [PATCH] simplify VM creation

---
 e2e/node_bootstrapper_test.go | 28 ++++++++++++++++++++++++++--
 e2e/types.go                  |  4 +---
 e2e/vmss.go                   | 11 +----------
 node-bootstrapper/app.go      |  3 ++-
 node-bootstrapper/main.go     |  1 -
 5 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/e2e/node_bootstrapper_test.go b/e2e/node_bootstrapper_test.go
index af7ce7baad6..3a2ec9bdb37 100644
--- a/e2e/node_bootstrapper_test.go
+++ b/e2e/node_bootstrapper_test.go
@@ -17,6 +17,8 @@ import (
 	"github.com/Azure/agentbaker/pkg/agent/datamodel"
 	nbcontractv1 "github.com/Azure/agentbaker/pkg/proto/nbcontract/v1"
 	"github.com/Azure/agentbakere2e/config"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v6"
 	"github.com/barkimedes/go-deepcopy"
 	"github.com/stretchr/testify/require"
 )
@@ -48,13 +50,35 @@ func Test_ubuntu2204NodeBootstrapper(t *testing.T) {
 			//NodeBootstrappingType: Scriptless,
 			Cluster: ClusterKubenet,
 			VHD:     config.VHDUbuntu2204Gen2Containerd,
+			VMConfigMutator: func(model *armcompute.VirtualMachineScaleSet) {
+				model.Properties.VirtualMachineProfile.OSProfile.CustomData = nil
+				model.Properties.VirtualMachineProfile.ExtensionProfile = &armcompute.VirtualMachineScaleSetExtensionProfile{
+					Extensions: []*armcompute.VirtualMachineScaleSetExtension{
+						{
+							Name: to.Ptr("vmssCSE"),
+							Properties: &armcompute.VirtualMachineScaleSetExtensionProperties{
+								Publisher:               to.Ptr("Microsoft.Azure.Extensions"),
+								Type:                    to.Ptr("CustomScript"),
+								TypeHandlerVersion:      to.Ptr("2.0"),
+								AutoUpgradeMinorVersion: to.Ptr(true),
+								Settings:                map[string]any{},
+								ProtectedSettings: map[string]any{
+									//"fileUris":         []string{},
+									"commandToExecute": CSENodeBootstrapper(ctx, t, cluster),
+									//"managedIdentity": map[string]any{
+									//	"objectId": "",
+									//},
+								},
+							},
+						},
+					},
+				}
+			},
 			LiveVMValidators: []*LiveVMValidator{
 				mobyComponentVersionValidator("containerd", getExpectedPackageVersions("containerd", "ubuntu", "r2204")[0], "apt"),
 				mobyComponentVersionValidator("runc", getExpectedPackageVersions("runc", "ubuntu", "r2204")[0], "apt"),
 				FileHasContentsValidator("/var/log/azure/node-bootstrapper.log", "node-bootstrapper finished successfully"),
 			},
-			CSEOverride:          CSENodeBootstrapper(ctx, t, cluster),
-			DisableCustomData:    true,
 			AKSNodeConfigMutator: func(config *nbcontractv1.Configuration) {},
 		},
 	})
diff --git a/e2e/types.go b/e2e/types.go
index a4923b2a6bd..90831f31dfa 100644
--- a/e2e/types.go
+++ b/e2e/types.go
@@ -140,9 +140,7 @@ type Config struct {
 
 	// LiveVMValidators is a slice of LiveVMValidator objects for performing any live VM validation
 	// specific to the scenario that isn't covered in the set of common validators run with all scenarios
-	LiveVMValidators  []*LiveVMValidator
-	CSEOverride       string
-	DisableCustomData bool
+	LiveVMValidators []*LiveVMValidator
 }
 
 // VMCommandOutputAsserterFn is a function which takes in stdout and stderr stream content
diff --git a/e2e/vmss.go b/e2e/vmss.go
index d544b6369fb..b403492190d 100644
--- a/e2e/vmss.go
+++ b/e2e/vmss.go
@@ -45,16 +45,7 @@ func createVMSS(ctx context.Context, t *testing.T, vmssName string, scenario *Sc
 		require.NoError(t, err)
 	}
 
-	cse := nodeBootstrapping.CSE
-	if scenario.CSEOverride != "" {
-		cse = scenario.CSEOverride
-	}
-	customData := nodeBootstrapping.CustomData
-	if scenario.DisableCustomData {
-		customData = ""
-	}
-
-	model := getBaseVMSSModel(vmssName, string(publicKeyBytes), customData, cse, cluster)
+	model := getBaseVMSSModel(vmssName, string(publicKeyBytes), nodeBootstrapping.CustomData, nodeBootstrapping.CSE, cluster)
 
 	isAzureCNI, err := cluster.IsAzureCNI()
 	require.NoError(t, err, "checking if cluster is using Azure CNI")
diff --git a/node-bootstrapper/app.go b/node-bootstrapper/app.go
index 19d60e0b42e..e36142a66ae 100644
--- a/node-bootstrapper/app.go
+++ b/node-bootstrapper/app.go
@@ -30,10 +30,11 @@ type ProvisionFlags struct {
 }
 
 func (a *App) Run(ctx context.Context, args []string) int {
+	slog.Info("node-bootstrapper started")
 	err := a.run(ctx, args)
 	exitCode := errToExitCode(err)
 	if exitCode == 0 {
-		slog.Info("node-bootstrapper finished")
+		slog.Info("node-bootstrapper finished successfully")
 	} else {
 		slog.Error("node-bootstrapper failed", "error", err)
 	}
diff --git a/node-bootstrapper/main.go b/node-bootstrapper/main.go
index 107636c8470..84fbbeb48a3 100644
--- a/node-bootstrapper/main.go
+++ b/node-bootstrapper/main.go
@@ -23,7 +23,6 @@ func main() {
 
 	logger := slog.New(slog.NewJSONHandler(logFile, nil))
 	slog.SetDefault(logger)
-	slog.Info("node-bootstrapper started")
 
 	app := App{cmdRunner: cmdRunner}
 	exitCode := app.Run(context.Background(), os.Args)