chore: coderabbit review comment changes

Signed-off-by: Tanisha goyal <tanishag@nvidia.com>

Author: tanishagoyal2

data-models/pkg/protos/health_event.pb.go

@@ -1604,10 +1604,12 @@ func TestProcessHealthEvents_StoreOnlyStrategy(t *testing.T) {
 			}
 
 			if tc.expectNodeConditions {
+				require.NotEmpty(t, nvsentinelConditions,
+					"Expected at least one NVSentinel node condition, but got none")
 				assert.Equal(t, tc.expectedConditionType, string(nvsentinelConditions[0].Type),
 					"Expected condition type %s, got %s", tc.expectedConditionType, nvsentinelConditions[0].Type)
 			} else {
-				assert.Equal(t, 0, len(nvsentinelConditions),
+				assert.Empty(t, nvsentinelConditions,
 					"Expected no NVSentinel node conditions for STORE_ONLY events, got %d", len(nvsentinelConditions))
 			}
 
@@ -1618,10 +1620,12 @@ func TestProcessHealthEvents_StoreOnlyStrategy(t *testing.T) {
 			require.NoError(t, err, "Failed to list events")
 
 			if tc.expectKubernetesEvents {
+				require.NotEmpty(t, events.Items,
+					"Expected at least one Kubernetes event, but got none")
 				assert.Equal(t, tc.expectedEventType, events.Items[0].Type,
 					"Expected event type %s, got %s", tc.expectedEventType, events.Items[0].Type)
 			} else {
-				assert.Equal(t, 0, len(events.Items),
+				assert.Empty(t, events.Items,
 					"Expected no Kubernetes events for STORE_ONLY events, got %d", len(events.Items))
 			}
 

platform-connectors/pkg/connectors/kubernetes/k8s_platform_connector_test.go

@@ -2272,56 +2272,41 @@ func waitForDaemonSetRollout(ctx context.Context, t *testing.T, client klient.Cl
 	t.Logf("DaemonSet %s/%s rollout completed successfully", NVSentinelNamespace, name)
 }
 
-// updateContainerProcessingStrategy updates the processing strategy argument for a specific container.
-func updateContainerProcessingStrategy(container *v1.Container) {
-	processingStrategyArg := "--processing-strategy=STORE_ONLY"
-
-	for j, arg := range container.Args {
-		if strings.HasPrefix(arg, "--processing-strategy=") {
-			container.Args[j] = processingStrategyArg
-			return
-		}
-
-		if arg == "--processing-strategy" && j+1 < len(container.Args) {
-			container.Args[j+1] = "STORE_ONLY"
-			return
-		}
-	}
-
-	container.Args = append(container.Args, processingStrategyArg)
-}
-
-// UpdateDaemonSetProcessingStrategy updates the daemonset to use STORE_ONLY processing strategy.
-func UpdateDaemonSetProcessingStrategy(ctx context.Context, t *testing.T,
-	client klient.Client, daemonsetName string, containerName string) (*appsv1.DaemonSet, error) {
+// UpdateDaemonSetArgs updates the daemonset with the specified arguments and complete the rollout.
+func UpdateDaemonSetArgs(ctx context.Context, t *testing.T,
+	client klient.Client, daemonsetName string, containerName string,
+	args map[string]string) error {
 	t.Helper()
 
-	t.Logf("Updating daemonset %s/%s to use STORE_ONLY processing strategy", NVSentinelNamespace, daemonsetName)
-
-	var originalDaemonSet *appsv1.DaemonSet
+	t.Logf("Updating daemonset %s/%s with args %v", NVSentinelNamespace, daemonsetName, args)
 
 	err := retry.RetryOnConflict(retry.DefaultRetry, func() error {
 		daemonSet := &appsv1.DaemonSet{}
 		if err := client.Resources().Get(ctx, daemonsetName, NVSentinelNamespace, daemonSet); err != nil {
 			return err
 		}
 
-		if originalDaemonSet == nil {
-			originalDaemonSet = daemonSet.DeepCopy()
-		}
+		containers := daemonSet.Spec.Template.Spec.Containers
+		containerFound := false
+
+		for i := range containers {
+			if containers[i].Name == containerName {
+				setArgsOnContainer(t, &containers[i], args)
+
+				containerFound = true
 
-		for i := range daemonSet.Spec.Template.Spec.Containers {
-			container := &daemonSet.Spec.Template.Spec.Containers[i]
-			if container.Name == containerName {
-				updateContainerProcessingStrategy(container)
 				break
 			}
 		}
 
+		if !containerFound {
+			return fmt.Errorf("container %q not found in daemonset %s/%s", containerName, NVSentinelNamespace, daemonsetName)
+		}
+
 		return client.Resources().Update(ctx, daemonSet)
 	})
 	if err != nil {
-		return nil, err
+		return err
 	}
 
 	t.Logf("Waiting for daemonset %s/%s rollout to complete", NVSentinelNamespace, daemonsetName)
@@ -2330,41 +2315,130 @@ func UpdateDaemonSetProcessingStrategy(ctx context.Context, t *testing.T,
 	t.Logf("Waiting 10 seconds for daemonset pods to start")
 	time.Sleep(10 * time.Second)
 
-	return originalDaemonSet, nil
+	return nil
 }
 
-func RestoreDaemonSet(ctx context.Context, t *testing.T, client klient.Client,
-	originalDaemonSet *appsv1.DaemonSet, daemonsetName string,
-) error {
+func RemoveDaemonSetArgs(ctx context.Context, t *testing.T, client klient.Client,
+	daemonsetName string,
+	containerName string, args map[string]string,
+) {
 	t.Helper()
 
-	if originalDaemonSet == nil {
-		t.Log("No original daemonset to restore, skipping")
-		return nil
-	}
-
-	t.Logf("Restoring daemonset %s/%s to original state", NVSentinelNamespace, daemonsetName)
+	t.Logf("Removing args %v from daemonset %s/%s", args, NVSentinelNamespace, daemonsetName)
 
 	err := retry.RetryOnConflict(retry.DefaultRetry, func() error {
 		daemonSet := &appsv1.DaemonSet{}
 		if err := client.Resources().Get(ctx, daemonsetName, NVSentinelNamespace, daemonSet); err != nil {
 			return err
 		}
 
-		daemonSet.Spec.Template.Spec.Containers = originalDaemonSet.Spec.Template.Spec.Containers
+		containers := daemonSet.Spec.Template.Spec.Containers
+		containerFound := false
+
+		for i := range containers {
+			if containers[i].Name == containerName {
+				removeArgsFromContainer(&containers[i], args)
+
+				containerFound = true
+
+				break
+			}
+		}
+
+		if !containerFound {
+			return fmt.Errorf("container %q not found in daemonset %s/%s", containerName, NVSentinelNamespace, daemonsetName)
+		}
 
 		return client.Resources().Update(ctx, daemonSet)
 	})
-	if err != nil {
-		return err
-	}
+	require.NoError(t, err, "failed to remove args from daemonset %s/%s", NVSentinelNamespace, daemonsetName)
 
 	t.Logf("Waiting for daemonset %s/%s rollout to complete after restoration", NVSentinelNamespace, daemonsetName)
 	waitForDaemonSetRollout(ctx, t, client, daemonsetName)
 
 	t.Log("DaemonSet restored successfully")
+}
 
-	return nil
+// tryUpdateExistingArg attempts to update an existing argument at position j.
+// Returns true if the argument was found and updated.
+func tryUpdateExistingArg(container *v1.Container, j int, flag, value string) bool {
+	existingArg := container.Args[j]
+
+	// Match --flag=value style
+	if strings.HasPrefix(existingArg, flag+"=") {
+		if value != "" {
+			container.Args[j] = flag + "=" + value
+		} else {
+			container.Args[j] = flag
+		}
+
+		return true
+	}
+
+	// Match --flag or --flag value style
+	if existingArg == flag {
+		if value != "" {
+			if j+1 < len(container.Args) && !strings.HasPrefix(container.Args[j+1], "-") {
+				container.Args[j+1] = value
+			} else {
+				container.Args = append(container.Args[:j+1], append([]string{value}, container.Args[j+1:]...)...)
+			}
+		}
+
+		return true
+	}
+
+	return false
+}
+
+func setArgsOnContainer(t *testing.T, container *v1.Container, args map[string]string) {
+	t.Helper()
+	t.Logf("Setting args %v on container %s", args, container.Name)
+
+	for flag, value := range args {
+		found := false
+
+		for j := 0; j < len(container.Args); j++ {
+			if tryUpdateExistingArg(container, j, flag, value) {
+				found = true
+				break
+			}
+		}
+
+		if !found {
+			if value != "" {
+				container.Args = append(container.Args, flag+"="+value)
+			} else {
+				container.Args = append(container.Args, flag)
+			}
+		}
+	}
+}
+
+func removeArgsFromContainer(container *v1.Container, args map[string]string) {
+	for flag := range args {
+		for j := 0; j < len(container.Args); j++ {
+			existingArg := container.Args[j]
+
+			// Match --flag=value style
+			if strings.HasPrefix(existingArg, flag+"=") {
+				container.Args = append(container.Args[:j], container.Args[j+1:]...)
+				break
+			}
+
+			// Match --flag or --flag value style
+
+			if existingArg == flag {
+				if j+1 < len(container.Args) && !strings.HasPrefix(container.Args[j+1], "-") {
+					container.Args = append(container.Args[:j], container.Args[j+2:]...)
+				} else {
+					container.Args = append(container.Args[:j], container.Args[j+1:]...)
+				}
+
+				break
+			}
+		}
+	}
 }
 
 func GetDaemonSetPodOnWorkerNode(ctx context.Context, t *testing.T, client klient.Client,

tests/helpers/kube.go

@@ -19,7 +19,6 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
-	appsv1 "k8s.io/api/apps/v1"
 	v1 "k8s.io/api/core/v1"
 	"sigs.k8s.io/e2e-framework/klient"
 )
@@ -32,15 +31,13 @@ const (
 
 // helper function to set up syslog health monitor and port forward to it
 func SetUpSyslogHealthMonitor(ctx context.Context, t *testing.T,
-	client klient.Client, updatedDaemonSet bool) (string, *appsv1.DaemonSet, *v1.Pod, chan struct{}) {
-	var originalDaemonSet *appsv1.DaemonSet
-
+	client klient.Client, args map[string]string) (string, *v1.Pod, chan struct{}) {
 	var err error
 
 	var syslogPod *v1.Pod
 
-	if updatedDaemonSet {
-		originalDaemonSet, err = UpdateDaemonSetProcessingStrategy(ctx, t, client, SyslogDaemonSetName, SyslogContainerName)
+	if args != nil {
+		err = UpdateDaemonSetArgs(ctx, t, client, SyslogDaemonSetName, SyslogContainerName, args)
 		require.NoError(t, err, "failed to update syslog health monitor processing strategy")
 	}
 
@@ -70,19 +67,18 @@ func SetUpSyslogHealthMonitor(ctx context.Context, t *testing.T,
 	err = SetNodeManagedByNVSentinel(ctx, client, testNodeName, false)
 	require.NoError(t, err, "failed to set ManagedByNVSentinel label")
 
-	return testNodeName, originalDaemonSet, syslogPod, stopChan
+	return testNodeName, syslogPod, stopChan
 }
 
 // helper function to roll back syslog health monitor daemonset and stop the port forward
 func TearDownSyslogHealthMonitor(ctx context.Context, t *testing.T, client klient.Client,
-	originalDaemonSet *appsv1.DaemonSet, nodeName string, stopChan chan struct{},
-	updatedDaemonSet bool, podName string) {
+	nodeName string, stopChan chan struct{},
+	args map[string]string, podName string) {
 	t.Log("Stopping port-forward")
 	close(stopChan)
 
-	if updatedDaemonSet {
-		err := RestoreDaemonSet(ctx, t, client, originalDaemonSet, "syslog-health-monitor-regular")
-		require.NoError(t, err, "failed to restore syslog health monitor daemon set")
+	if args != nil {
+		RemoveDaemonSetArgs(ctx, t, client, SyslogDaemonSetName, SyslogContainerName, args)
 	}
 
 	t.Logf("Restarting syslog-health-monitor pod %s to clear conditions", podName)