test(diff): derive cancellation checkpoints

Signed-off-by: Tjark Gunnar Rasche <trasche@nvidia.com>

Author: tjrasche

pkg/client/v1/diff.go

@@ -213,6 +213,6 @@ func copySnapshotDiffString(value *string) *string {
 	if value == nil {
 		return nil
 	}
-	copy := *value
-	return &copy
+	copied := *value
+	return &copied
 }

pkg/client/v1/diff_test.go

@@ -255,6 +255,19 @@ func TestDiffSnapshots_MidTraversalContextCancellation(t *testing.T) {
 	}
 	baseline := diffTestSnapshot(baselineData, nil, nil)
 	target := diffTestSnapshot(targetData, nil, nil)
+	probeCtx := &snapshotDiffCountingContext{Context: t.Context()}
+	probeResult, err := client.DiffSnapshots(probeCtx, baseline, target, aicr.SnapshotDiffOptions{})
+	if err != nil {
+		t.Fatalf("DiffSnapshots() probe error = %v", err)
+	}
+	// Exclude the facade's one mapping checkpoint per change, then cancel
+	// halfway through the comparison's final summary traversal. This proves
+	// the internal partial result is discarded rather than failing in mapping.
+	comparisonChecks := probeCtx.checks - len(probeResult.Changes)
+	cancelAt := comparisonChecks - probeResult.Summary.Total/2
+	if cancelAt <= 0 || cancelAt >= comparisonChecks {
+		t.Fatalf("derived cancellation checkpoint = %d, comparison checks = %d", cancelAt, comparisonChecks)
+	}
 
 	tests := []struct {
 		name     string
@@ -267,10 +280,6 @@ func TestDiffSnapshots_MidTraversalContextCancellation(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			// The first 212 checkpoints index and sort this fixture. Canceling
-			// later proves that an in-progress result with accumulated changes
-			// is discarded rather than escaping through the facade.
-			const cancelAt = 230
 			ctx := newSnapshotDiffCheckpointContext(t.Context(), cancelAt, tt.cause)
 			result, err := client.DiffSnapshots(ctx, baseline, target, aicr.SnapshotDiffOptions{})
 			if result != nil {
@@ -394,6 +403,16 @@ type snapshotDiffCheckpointContext struct {
 	closed   bool
 }
 
+type snapshotDiffCountingContext struct {
+	context.Context
+	checks int
+}
+
+func (c *snapshotDiffCountingContext) Err() error {
+	c.checks++
+	return c.Context.Err()
+}
+
 func newSnapshotDiffCheckpointContext(
 	parent context.Context,
 	cancelAt int,

pkg/diff/diff_test.go

@@ -700,22 +700,31 @@ func TestSnapshotsWithContext_MidTraversalCancellation(t *testing.T) {
 	}
 	baseline := makeSnapshot(makeMeasurement(measurement.TypeK8s, makeSubtype("server", baselineData)))
 	target := makeSnapshot(makeMeasurement(measurement.TypeK8s, makeSubtype("server", targetData)))
+	probeCtx := &countingContext{Context: t.Context()}
+	probeResult, err := SnapshotsWithContext(probeCtx, baseline, target)
+	if err != nil {
+		t.Fatalf("SnapshotsWithContext() probe error = %v", err)
+	}
+	// The final summary traversal checks the context once per accumulated
+	// change. Cancel halfway through it without depending on the number of
+	// checkpoints used by earlier comparison stages.
+	cancelAt := probeCtx.checks - probeResult.Summary.Total/2
+	if cancelAt <= 0 || cancelAt >= probeCtx.checks {
+		t.Fatalf("derived cancellation checkpoint = %d, probe checks = %d", cancelAt, probeCtx.checks)
+	}
 
 	tests := []struct {
 		name     string
 		cause    error
 		wantCode aicrerrors.ErrorCode
-		cancelAt int
 	}{
-		// The first 212 checkpoints index and sort this fixture. Canceling
-		// later proves that accumulated changes are discarded with the result.
-		{name: "canceled", cause: context.Canceled, wantCode: aicrerrors.ErrCodeCanceled, cancelAt: 230},
-		{name: "deadline", cause: context.DeadlineExceeded, wantCode: aicrerrors.ErrCodeTimeout, cancelAt: 230},
+		{name: "canceled", cause: context.Canceled, wantCode: aicrerrors.ErrCodeCanceled},
+		{name: "deadline", cause: context.DeadlineExceeded, wantCode: aicrerrors.ErrCodeTimeout},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			ctx := newCheckpointContext(t.Context(), tt.cancelAt, tt.cause)
+			ctx := newCheckpointContext(t.Context(), cancelAt, tt.cause)
 			result, err := SnapshotsWithContext(ctx, baseline, target)
 			if result != nil {
 				t.Fatalf("SnapshotsWithContext() result = %#v, want nil after cancellation", result)
@@ -726,8 +735,8 @@ func TestSnapshotsWithContext_MidTraversalCancellation(t *testing.T) {
 			if !stderrors.Is(err, tt.cause) {
 				t.Errorf("SnapshotsWithContext() error = %v, want cause %v", err, tt.cause)
 			}
-			if ctx.checks < tt.cancelAt {
-				t.Errorf("context checks = %d, want at least %d to prove traversal began", ctx.checks, tt.cancelAt)
+			if ctx.checks < cancelAt {
+				t.Errorf("context checks = %d, want at least %d to prove traversal began", ctx.checks, cancelAt)
 			}
 		})
 	}
@@ -770,6 +779,16 @@ type checkpointContext struct {
 	closed   bool
 }
 
+type countingContext struct {
+	context.Context
+	checks int
+}
+
+func (c *countingContext) Err() error {
+	c.checks++
+	return c.Context.Err()
+}
+
 func newCheckpointContext(parent context.Context, cancelAt int, cause error) *checkpointContext {
 	return &checkpointContext{
 		Context:  parent,