update test helpers

Azure · Feb 26, 2025 · 8bdfb30 · 8bdfb30
1 parent e5ee2d8
commit 8bdfb30
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 31 deletions.
diff --git a/e2e/kube.go b/e2e/kube.go
@@ -452,13 +452,7 @@ func getClusterSubnetID(ctx context.Context, mcResourceGroupName string, t *test
 }
 
 func podHTTPServerLinux(s *Scenario) *corev1.Pod {
-	image := "mcr.microsoft.com/cbl-mariner/busybox:2.0"
-	secretName := ""
-	if s.Tags.Airgap {
-		image = fmt.Sprintf("%s.azurecr.io/cbl-mariner/busybox:2.0", config.GetPrivateACRName(s.Tags.NonAnonymousACR))
-		secretName = config.Config.ACRSecretName
-	}
-	return &corev1.Pod{
+	pod := &corev1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      fmt.Sprintf("%s-test-pod", s.Runtime.KubeNodeName),
 			Namespace: "default",
@@ -467,7 +461,7 @@ func podHTTPServerLinux(s *Scenario) *corev1.Pod {
 			Containers: []corev1.Container{
 				{
 					Name:  "mariner",
-					Image: image,
+					Image: "mcr.microsoft.com/cbl-mariner/busybox:2.0",
 					Ports: []corev1.ContainerPort{
 						{
 							ContainerPort: 80,
@@ -499,13 +493,13 @@ func podHTTPServerLinux(s *Scenario) *corev1.Pod {
 			NodeSelector: map[string]string{
 				"kubernetes.io/hostname": s.Runtime.KubeNodeName,
 			},
-			ImagePullSecrets: []corev1.LocalObjectReference{
-				{
-					Name: secretName,
-				},
-			},
 		},
 	}
+	if s.Tags.Airgap {
+		pod.Spec.Containers[0].Image = fmt.Sprintf("%s.azurecr.io/cbl-mariner/busybox:2.0", config.GetPrivateACRName(s.Tags.NonAnonymousACR))
+		pod.Spec.ImagePullSecrets = []corev1.LocalObjectReference{{Name: config.Config.ACRSecretName}}
+	}
+	return pod
 }
 
 func podHTTPServerWindows(s *Scenario) *corev1.Pod {

diff --git a/e2e/scenario_helpers_test.go b/e2e/scenario_helpers_test.go
@@ -146,8 +146,6 @@ func prepareAKSNode(ctx context.Context, s *Scenario) {
 	s.T.Logf("vmss %s creation succeeded", s.Runtime.VMSSName)
 
 	s.Runtime.KubeNodeName = s.Runtime.Cluster.Kube.WaitUntilNodeReady(ctx, s.T, s.Runtime.VMSSName)
-	s.T.Logf("node %s is ready", s.Runtime.VMSSName)
-
 }
 
 func maybeSkipScenario(ctx context.Context, t *testing.T, s *Scenario) {

diff --git a/e2e/validators.go b/e2e/validators.go
@@ -451,6 +451,7 @@ func ValidateAMDGPU(ctx context.Context, s *Scenario) {
 	assert.Contains(s.T, execResult.stdout.String(), "amdgpu", "expected to see amdgpu kernel module managing a PCI device, but did not")
 
 	ensurePod(ctx, s, podEnableAMDGPUResource(s))
+	s.T.Logf("waiting for AMD GPU to be available")
 	waitUntilResourceAvailable(ctx, s, "amd.com/gpu")
 	//ensureJob(ctx, s, jobAMDGPUWorkload(s))
 }
diff --git a/e2e/vmss.go b/e2e/vmss.go
@@ -14,6 +14,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -25,6 +26,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v6"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/crypto/ssh"
 )
@@ -35,14 +37,6 @@ const (
 )
 
 func createVMSS(ctx context.Context, s *Scenario) *armcompute.VirtualMachineScaleSet {
-	defer func() {
-		var err error
-		s.Runtime.VMPrivateIP, err = getVMPrivateIPAddress(ctx, s)
-		require.NoError(s.T, err, "failed to get VM private IP address")
-
-		uploadSSHKey(ctx, s)
-		logSSHInstructions(s)
-	}()
 
 	cluster := s.Runtime.Cluster
 	var nodeBootstrapping *datamodel.NodeBootstrapping
@@ -81,11 +75,15 @@ func createVMSS(ctx context.Context, s *Scenario) *armcompute.VirtualMachineScal
 	}
 
 	s.PrepareVMSSModel(ctx, s.T, &model)
-
 	vmss, err := config.Azure.CreateVMSSWithRetry(ctx, s.T, *cluster.Model.Properties.NodeResourceGroup, s.Runtime.VMSSName, model)
 	s.T.Cleanup(func() {
 		cleanupVMSS(ctx, s)
 	})
+	var ipErr error
+	s.Runtime.VMPrivateIP, ipErr = getVMPrivateIPAddress(ctx, s)
+	assert.NoError(s.T, ipErr, "failed to get VM private IP address")
+	uploadSSHKey(ctx, s)
+	logSSHInstructions(s)
 	skipTestIfSKUNotAvailableErr(s.T, err)
 	// fail test, but continue to extract debug information
 	require.NoError(s.T, err, "create vmss %q, check %s for vm logs", s.Runtime.VMSSName, testDir(s.T))
@@ -129,14 +127,23 @@ func extractLogsFromVMLinux(ctx context.Context, s *Scenario) {
 	}
 
 	var logFiles = map[string]string{}
+	wg := sync.WaitGroup{}
+	lock := sync.Mutex{}
 	for file, sourceCmd := range commandList {
-		execResult, err := execBashCommandOnVM(ctx, s, sourceCmd)
-		if err != nil {
-			s.T.Logf("error executing %s: %s", sourceCmd, err)
-			continue
-		}
-		logFiles[file] = execResult.String()
+		wg.Add(1)
+		go func(file, sourceCmd string) {
+			defer wg.Done()
+			execResult, err := execBashCommandOnVM(ctx, s, sourceCmd)
+			if err != nil {
+				s.T.Logf("error executing %s: %s", sourceCmd, err)
+				return
+			}
+			lock.Lock()
+			logFiles[file] = execResult.String()
+			lock.Unlock()
+		}(file, sourceCmd)
 	}
+	wg.Wait()
 	err := dumpFileMapToDir(s.T, logFiles)
 	require.NoError(s.T, err)
 }
@@ -505,7 +512,7 @@ func getBaseVMSSModel(s *Scenario, customData, cseCmd string) armcompute.Virtual
 					Properties: &armcompute.VirtualMachineScaleSetExtensionProperties{
 						Publisher:               to.Ptr("Microsoft.Azure.Extensions"),
 						Type:                    to.Ptr("CustomScript"),
-						TypeHandlerVersion:      to.Ptr("2.0"),
+						TypeHandlerVersion:      to.Ptr("2.1"),
 						AutoUpgradeMinorVersion: to.Ptr(true),
 						Settings:                map[string]interface{}{},
 						ProtectedSettings: map[string]interface{}{