Kubernetes tests: Let the benchmark metric recorder be overridden via context

PiperOrigin-RevId: 700842559

test/kubernetes/benchmetric/BUILD

@@ -2,7 +2,7 @@ load("//tools:defs.bzl", "go_library")
 
 package(
     default_applicable_licenses = ["//:license"],
-    default_visibility = ["//test/kubernetes:__subpackages__"],
+    default_visibility = ["//:sandbox"],
     licenses = ["notice"],
 )
 

test/kubernetes/benchmetric/benchmetric.go

@@ -178,8 +178,20 @@ var (
 	recorderOnce sync.Once
 )
 
+type recorderContextKeyType int
+
+const recorderContextKey recorderContextKeyType = iota
+
+// WithRecorder returns a context with the given `Recorder`.
+func WithRecorder(ctx context.Context, recorder Recorder) context.Context {
+	return context.WithValue(ctx, recorderContextKey, recorder)
+}
+
 // GetRecorder returns the benchmark's `Recorder` singleton.
 func GetRecorder(ctx context.Context) (Recorder, error) {
+	if ctx.Value(recorderContextKey) != nil {
+		return ctx.Value(recorderContextKey).(Recorder), nil
+	}
 	recorderOnce.Do(func() {
 		recorder, recorderErr = recorderFn(ctx)
 	})