fix: allow user to only set one field for ng or cz (#211) (#213)

Signed-off-by: Amber Xue <ambermingxin@nvidia.com>
Signed-off-by: Jingxiang Zhang <jingzhang@nvidia.com>
Co-authored-by: ambermingxin <ambermingxin@nvidia.com>

Author: jingxiang-z

cmd/fleetint/enroll.go

@@ -96,9 +96,6 @@ func enrollCommand(cliContext *cli.Context) error {
 	if err != nil {
 		return err
 	}
-	if err := validateReservedPairMetadata(nodeGroup, computeZone); err != nil {
-		return err
-	}
 	metadata := &enrollment.EnrollMetadata{
 		NodeGroup:   nodeGroup,
 		ComputeZone: computeZone,
@@ -159,20 +156,3 @@ func validatedOptionalMetadataFlagValue(cliContext *cli.Context, name, fieldName
 	}
 	return &value, nil
 }
-
-func validateReservedPairMetadata(nodeGroup, computeZone *string) error {
-	nodeGroupSet := nodeGroup != nil
-	computeZoneSet := computeZone != nil
-
-	if !nodeGroupSet && !computeZoneSet {
-		return nil
-	}
-	if nodeGroupSet && computeZoneSet {
-		nodeGroupEmpty := *nodeGroup == ""
-		computeZoneEmpty := *computeZone == ""
-		if nodeGroupEmpty == computeZoneEmpty {
-			return nil
-		}
-	}
-	return fmt.Errorf("--node-group and --compute-zone must be both omitted, both empty, or both non-empty")
-}

cmd/fleetint/enroll_test.go

@@ -340,49 +340,83 @@ func TestEnrollCommandRejectsOverlongMetadataNames(t *testing.T) {
 	require.ErrorContains(t, err, "Node group name must be 255 characters or fewer")
 }
 
-func TestValidateReservedPairMetadata(t *testing.T) {
+func TestValidatedOptionalMetadataFlagValueSupportsIndependentReservedUpdates(t *testing.T) {
 	strPtr := func(v string) *string { return &v }
 	tests := []struct {
-		name        string
-		nodeGroup   *string
-		computeZone *string
-		wantErr     bool
+		name              string
+		nodeGroupArg      *string
+		computeZoneArg    *string
+		expectNodeGroup   *string
+		expectComputeZone *string
 	}{
-		{name: "both omitted", nodeGroup: nil, computeZone: nil, wantErr: false},
-		{name: "both empty", nodeGroup: strPtr(""), computeZone: strPtr(""), wantErr: false},
-		{name: "both non-empty", nodeGroup: strPtr("ng-a"), computeZone: strPtr("cz-a"), wantErr: false},
-		{name: "node-group only", nodeGroup: strPtr("ng-a"), computeZone: nil, wantErr: true},
-		{name: "compute-zone only", nodeGroup: nil, computeZone: strPtr("cz-a"), wantErr: true},
-		{name: "node-group empty compute-zone non-empty", nodeGroup: strPtr(""), computeZone: strPtr("cz-a"), wantErr: true},
-		{name: "node-group non-empty compute-zone empty", nodeGroup: strPtr("ng-a"), computeZone: strPtr(""), wantErr: true},
+		{
+			name:              "clear compute-zone and set node-group",
+			nodeGroupArg:      strPtr("xx"),
+			computeZoneArg:    strPtr(""),
+			expectNodeGroup:   strPtr("xx"),
+			expectComputeZone: strPtr(""),
+		},
+		{
+			name:              "set compute-zone and clear node-group",
+			nodeGroupArg:      strPtr(""),
+			computeZoneArg:    strPtr("xx"),
+			expectNodeGroup:   strPtr(""),
+			expectComputeZone: strPtr("xx"),
+		},
+		{
+			name:              "set compute-zone and omit node-group",
+			nodeGroupArg:      nil,
+			computeZoneArg:    strPtr("xx"),
+			expectNodeGroup:   nil,
+			expectComputeZone: strPtr("xx"),
+		},
+		{
+			name:              "set node-group and omit compute-zone",
+			nodeGroupArg:      strPtr("xx"),
+			computeZoneArg:    nil,
+			expectNodeGroup:   strPtr("xx"),
+			expectComputeZone: nil,
+		},
+		{
+			name:              "clear compute-zone and omit node-group",
+			nodeGroupArg:      nil,
+			computeZoneArg:    strPtr(""),
+			expectNodeGroup:   nil,
+			expectComputeZone: strPtr(""),
+		},
+		{
+			name:              "clear node-group and omit compute-zone",
+			nodeGroupArg:      strPtr(""),
+			computeZoneArg:    nil,
+			expectNodeGroup:   strPtr(""),
+			expectComputeZone: nil,
+		},
 	}
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			err := validateReservedPairMetadata(tc.nodeGroup, tc.computeZone)
-			if tc.wantErr {
-				require.ErrorContains(t, err, "--node-group and --compute-zone must be both omitted, both empty, or both non-empty")
-				return
+			flagSet := flag.NewFlagSet("enroll", flag.ContinueOnError)
+			flagSet.String("node-group", "", "")
+			flagSet.String("compute-zone", "", "")
+			if tc.nodeGroupArg != nil {
+				require.NoError(t, flagSet.Set("node-group", *tc.nodeGroupArg))
 			}
+			if tc.computeZoneArg != nil {
+				require.NoError(t, flagSet.Set("compute-zone", *tc.computeZoneArg))
+			}
+			cliContext := cli.NewContext(cli.NewApp(), flagSet, nil)
+
+			nodeGroup, err := validatedOptionalMetadataFlagValue(cliContext, "node-group", "Node group")
 			require.NoError(t, err)
+			computeZone, err := validatedOptionalMetadataFlagValue(cliContext, "compute-zone", "Compute zone")
+			require.NoError(t, err)
+
+			assert.Equal(t, tc.expectNodeGroup, nodeGroup)
+			assert.Equal(t, tc.expectComputeZone, computeZone)
 		})
 	}
 }
 
-func TestEnrollCommandRejectsMixedReservedPairValues(t *testing.T) {
-	app := App()
-	app.Writer = &bytes.Buffer{}
-
-	err := app.Run([]string{
-		"fleetint", "enroll",
-		"--endpoint", "https://example.com",
-		"--token", "token",
-		"--node-group", "",
-		"--compute-zone", "cz-a",
-	})
-	require.ErrorContains(t, err, "--node-group and --compute-zone must be both omitted, both empty, or both non-empty")
-}
-
 func TestEnrollCommandRejectsReservedUnassignedName(t *testing.T) {
 	app := App()
 	app.Writer = &bytes.Buffer{}

deployments/helm/fleet-intelligence-agent/values.yaml

@@ -61,13 +61,10 @@ enroll:
   tokenSecretName: ""
   tokenSecretKey: "token"
   tokenValue: ""
-  # Optional enrollment metadata:
-  # - omit key to preserve existing stored value
-  # - set non-empty string to overwrite
-  # - set empty string ("") to clear
-  #
-  # nodeGroup: "prod-a"
-  # computeZone: "us-east-1c"
+  # Optional enrollment metadata. Omit these keys to preserve stored values.
+  # Set to "" explicitly to clear stored values.
+  # nodeGroup: ""
+  # computeZone: ""
 
 ports:
   http: 15133

internal/agentstate/sqlite_test.go

@@ -23,6 +23,7 @@ import (
 	"testing"
 	"time"
 
+	pkgmetadata "github.com/NVIDIA/fleet-intelligence-sdk/pkg/metadata"
 	"github.com/NVIDIA/fleet-intelligence-sdk/pkg/sqlite"
 	"github.com/stretchr/testify/require"
 )
@@ -84,6 +85,15 @@ func TestSQLiteStateRoundTrip(t *testing.T) {
 	require.True(t, ok)
 	require.Equal(t, "group-a", value)
 
+	stateFile, err := state.stateFileFn()
+	require.NoError(t, err)
+	db, err := sqlite.Open(stateFile, sqlite.WithReadOnly(true))
+	require.NoError(t, err)
+	defer db.Close()
+	value, err = pkgmetadata.ReadMetadata(ctx, db, "node_group")
+	require.NoError(t, err)
+	require.Equal(t, "group-a", value)
+
 	value, ok, err = state.GetComputeZone(ctx)
 	require.NoError(t, err)
 	require.True(t, ok)

internal/agentstate/state.go

@@ -25,7 +25,7 @@ const (
 	MetadataKeyBackendBaseURL = "backend_base_url"
 	MetadataKeySAKToken       = "sak_token"
 	MetadataKeyEnrolledAt     = "enrolled_at"
-	MetadataKeyNodeGroup      = "nodegroup"
+	MetadataKeyNodeGroup      = "node_group"
 	MetadataKeyComputeZone    = "compute_zone"
 )
 

internal/enrollment/enrollment.go

@@ -175,7 +175,7 @@ func storeConfigInMetadata(ctx context.Context, baseURL, jwtToken, sakToken stri
 	}
 	if metadata.NodeGroup != nil {
 		if err := pkgmetadata.SetMetadata(ctx, dbRW, agentstate.MetadataKeyNodeGroup, *metadata.NodeGroup); err != nil {
-			return fmt.Errorf("failed to set nodegroup: %w", err)
+			return fmt.Errorf("failed to set node_group: %w", err)
 		}
 	}
 	if metadata.ComputeZone != nil {

internal/exporter/exporter.go

@@ -196,7 +196,7 @@ func (e *healthExporter) populateOptionalResourceMetadata(ctx context.Context, d
 
 	nodeGroup, err := pkgmetadata.ReadMetadata(ctx, e.options.dbRO, agentstate.MetadataKeyNodeGroup)
 	if err != nil {
-		log.Logger.Debugw("nodegroup metadata not available for telemetry resource", "error", err)
+		log.Logger.Debugw("node_group metadata not available for telemetry resource", "error", err)
 	} else {
 		data.NodeGroup = nodeGroup
 	}

internal/inventory/sink/backend.go

@@ -83,7 +83,7 @@ func (s *backendSink) Export(ctx context.Context, snap *inventory.Snapshot) erro
 	req := mapper.ToNodeUpsertRequest(snap)
 	nodeGroup, ok, err := s.state.GetNodeGroup(ctx)
 	if err != nil {
-		log.Logger.Warnw("inventory export continuing without nodegroup metadata", "error", err)
+		log.Logger.Warnw("inventory export continuing without node_group metadata", "error", err)
 	} else if ok {
 		req.NodeGroup = nodeGroup
 	}