chore: more changes

Signed-off-by: Ajay Mishra <[email protected]>

Author: XRFXLP

tests/data/basic-matching-configmap.yaml

@@ -19,6 +19,10 @@ metadata:
   namespace: nvsentinel
 data:
   config.toml: |
+    [circuitBreaker]
+      percentage = 50
+      duration = "5m"
+    
     [[rule-sets]]
       name = "Basic-Match-Rule"
       [[rule-sets.match.all]]

tests/data/managed-by-nvsentinel-configmap.yaml

@@ -19,6 +19,10 @@ metadata:
   namespace: nvsentinel
 data:
   config.toml: |
+    [circuitBreaker]
+      percentage = 50
+      duration = "5m"
+    
     [[rule-sets]]
       name = "ManagedByNVSentinel-Rule"
       [[rule-sets.match.all]]

tests/go.mod

@@ -5,6 +5,7 @@ go 1.25
 toolchain go1.25.3
 
 require (
+	github.com/BurntSushi/toml v1.5.0
 	github.com/go-logr/logr v1.4.3
 	github.com/nvidia/nvsentinel/commons v0.0.0
 	github.com/prometheus/client_golang v1.23.2

tests/go.sum

@@ -1,3 +1,5 @@
+github.com/BurntSushi/toml v1.5.0 h1:W5quZX/G/csjUnuI8SUYlsHs9M38FC7znL0lIO+DvMg=
+github.com/BurntSushi/toml v1.5.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho=
 github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio=
 github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=

tests/helpers/fault_quarantine.go

@@ -37,6 +37,17 @@ const (
 	CELKeyConfigMapBackup
 )
 
+type faultQuarantineConfig struct {
+	LabelPrefix    string               `toml:"label-prefix"`
+	CircuitBreaker circuitBreakerConfig `toml:"circuitBreaker"`
+	RuleSets       []map[string]any     `toml:"rule-sets"`
+}
+
+type circuitBreakerConfig struct {
+	Percentage int    `toml:"percentage"`
+	Duration   string `toml:"duration"`
+}
+
 type QuarantineTestContext struct {
 	NodeName        string
 	ConfigMapBackup []byte
@@ -94,16 +105,11 @@ func SetupQuarantineTest(ctx context.Context, t *testing.T, c *envconf.Config, c
 
 // QuarantineSetupOptions provides options for setting up quarantine tests.
 type QuarantineSetupOptions struct {
-	// CircuitBreakerPercentage sets the CB percentage threshold (0 to skip)
 	CircuitBreakerPercentage int
-	// CircuitBreakerDuration sets the CB duration (empty to skip)
-	CircuitBreakerDuration string
-	// CircuitBreakerState sets the initial CB state (empty to skip)
-	CircuitBreakerState string
-	// DryRun sets dry-run mode (nil to skip, otherwise pointer to bool value)
-	DryRun *bool
-	// SkipRestart skips the deployment restart (useful when chaining operations)
-	SkipRestart bool
+	CircuitBreakerDuration   string
+	CircuitBreakerState      string
+	DryRun                   *bool
+	SkipRestart              bool
 }
 
 // SetupQuarantineTestWithOptions sets up a quarantine test with additional configuration options.
@@ -118,36 +124,31 @@ func SetupQuarantineTestWithOptions(ctx context.Context, t *testing.T, c *envcon
 	testCtx := &QuarantineTestContext{}
 	var originalDeployment *appsv1.Deployment
 
-	t.Log("Backing up current fault-quarantine configmap")
 	backupData, err := BackupConfigMap(ctx, client, "fault-quarantine", NVSentinelNamespace)
 	require.NoError(t, err)
-	t.Log("Backup created in memory")
 	testCtx.ConfigMapBackup = backupData
 
 	t.Logf("Applying test configmap: %s", configMapPath)
 	err = createConfigMapFromFilePath(ctx, client, configMapPath, "fault-quarantine", NVSentinelNamespace)
 	require.NoError(t, err)
 
-	argUpdates := make(map[string]string)
 	if opts != nil {
 		if opts.CircuitBreakerPercentage > 0 {
-			t.Logf("Will set circuit breaker threshold: %d%%, duration: %s",
+			t.Logf("Updating circuit breaker in ConfigMap: %d%%, duration: %s",
 				opts.CircuitBreakerPercentage, opts.CircuitBreakerDuration)
-			argUpdates["--circuit-breaker-percentage="] = fmt.Sprintf("--circuit-breaker-percentage=%d", opts.CircuitBreakerPercentage)
-			argUpdates["--circuit-breaker-duration="] = fmt.Sprintf("--circuit-breaker-duration=%s", opts.CircuitBreakerDuration)
+			err = updateCircuitBreakerConfigInConfigMap(ctx, t, client, opts.CircuitBreakerPercentage, opts.CircuitBreakerDuration)
+			require.NoError(t, err)
 		}
 		if opts.DryRun != nil {
 			t.Logf("Will set dry-run mode to: %v", *opts.DryRun)
-			argUpdates["--dry-run="] = fmt.Sprintf("--dry-run=%v", *opts.DryRun)
+			argUpdates := map[string]string{
+				"--dry-run=": fmt.Sprintf("--dry-run=%v", *opts.DryRun),
+			}
+			originalDeployment = modifyFaultQuarantineDeploymentArgs(ctx, t, client, argUpdates)
+		}
+		if opts.CircuitBreakerState != "" {
+			updateCircuitBreakerStateConfigMap(ctx, t, client, opts.CircuitBreakerState)
 		}
-	}
-
-	if len(argUpdates) > 0 {
-		originalDeployment = modifyFaultQuarantineDeploymentArgs(ctx, t, client, argUpdates)
-	}
-
-	if opts != nil && opts.CircuitBreakerState != "" {
-		updateCircuitBreakerStateConfigMap(ctx, t, client, opts.CircuitBreakerState)
 	}
 
 	if opts == nil || !opts.SkipRestart {
@@ -479,3 +480,22 @@ func updateCircuitBreakerStateConfigMap(ctx context.Context, t *testing.T, clien
 	err := client.Resources().Create(ctx, cm)
 	require.NoError(t, err, "failed to create CB state configmap")
 }
+
+// updateCircuitBreakerConfigInConfigMap updates the circuit breaker percentage and duration in the fault-quarantine ConfigMap.
+func updateCircuitBreakerConfigInConfigMap(ctx context.Context, t *testing.T, client klient.Client, percentage int, duration string) error {
+	t.Helper()
+
+	err := UpdateConfigMapTOMLField(ctx, client, "fault-quarantine", NVSentinelNamespace, "config.toml",
+		func(cfg *faultQuarantineConfig) error {
+			cfg.CircuitBreaker.Percentage = percentage
+			cfg.CircuitBreaker.Duration = duration
+			return nil
+		})
+
+	if err != nil {
+		return fmt.Errorf("failed to update circuit breaker config: %w", err)
+	}
+
+	t.Logf("Updated circuit breaker config in ConfigMap: percentage=%d, duration=%s", percentage, duration)
+	return nil
+}

tests/helpers/kube.go

@@ -29,6 +29,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/BurntSushi/toml"
 	"github.com/nvidia/nvsentinel/commons/pkg/statemanager"
 	"github.com/stretchr/testify/require"
 	"go.yaml.in/yaml/v2"
@@ -726,6 +727,60 @@ func BackupConfigMap(ctx context.Context, c klient.Client, name, namespace strin
 	return yamlData, nil
 }
 
+// UpdateConfigMapTOMLField updates a TOML ConfigMap using proper TOML parsing.
+// The modifier function receives the parsed config and can modify it before it's written back.
+//
+// Example usage:
+//
+//	err := UpdateConfigMapTOMLField(ctx, client, "fault-quarantine", "nvsentinel", "config.toml",
+//	    func(cfg *map[string]interface{}) error {
+//	        if cb, ok := (*cfg)["circuitBreaker"].(map[string]interface{}); ok {
+//	            cb["percentage"] = 40
+//	            cb["duration"] = "10m"
+//	        }
+//	        return nil
+//	    })
+func UpdateConfigMapTOMLField[T any](ctx context.Context, c klient.Client, name, namespace, tomlKey string,
+	modifier func(*T) error) error {
+
+	cm := &v1.ConfigMap{}
+	err := c.Resources().Get(ctx, name, namespace, cm)
+	if err != nil {
+		return fmt.Errorf("failed to get configmap %s/%s: %w", namespace, name, err)
+	}
+
+	if cm.Data == nil {
+		return fmt.Errorf("configmap %s/%s has no data", namespace, name)
+	}
+
+	tomlData, exists := cm.Data[tomlKey]
+	if !exists {
+		return fmt.Errorf("configmap %s/%s missing key %s", namespace, name, tomlKey)
+	}
+
+	var cfg T
+	if err := toml.Unmarshal([]byte(tomlData), &cfg); err != nil {
+		return fmt.Errorf("failed to unmarshal TOML: %w", err)
+	}
+
+	if err := modifier(&cfg); err != nil {
+		return fmt.Errorf("modifier failed: %w", err)
+	}
+
+	var buf strings.Builder
+	encoder := toml.NewEncoder(&buf)
+	if err := encoder.Encode(&cfg); err != nil {
+		return fmt.Errorf("failed to marshal TOML: %w", err)
+	}
+
+	cm.Data[tomlKey] = buf.String()
+	if err := c.Resources().Update(ctx, cm); err != nil {
+		return fmt.Errorf("failed to update configmap %s/%s: %w", namespace, name, err)
+	}
+
+	return nil
+}
+
 // WaitForDeploymentRollout waits for a deployment rollout to complete.
 // This checks the same conditions as `kubectl rollout status` and verifies at least one pod is ready.
 func WaitForDeploymentRollout(ctx context.Context, t *testing.T, c klient.Client, name, namespace string) {