Internal change (diffbased).

PiperOrigin-RevId: 694595149

Author: EtiennePerot

test/kubernetes/benchmarks/BUILD

@@ -8,6 +8,7 @@ package(
 go_test(
     name = "abslbuild_test",
     srcs = ["abslbuild_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",
@@ -26,6 +27,7 @@ go_test(
 go_test(
     name = "startup_test",
     srcs = ["startup_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",
@@ -43,6 +45,7 @@ go_test(
 go_test(
     name = "redis_test",
     srcs = ["redis_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",
@@ -61,6 +64,7 @@ go_test(
 go_test(
     name = "ruby_dev_test",
     srcs = ["ruby_dev_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",
@@ -80,6 +84,7 @@ go_test(
 go_test(
     name = "ffmpeg_test",
     srcs = ["ffmpeg_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",
@@ -98,6 +103,7 @@ go_test(
 go_test(
     name = "grpc_test",
     srcs = ["grpc_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",
@@ -116,6 +122,7 @@ go_test(
 go_test(
     name = "nginx_test",
     srcs = ["nginx_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",
@@ -135,6 +142,7 @@ go_test(
 go_test(
     name = "postgresql_test",
     srcs = ["postgresql_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",
@@ -153,6 +161,7 @@ go_test(
 go_test(
     name = "tensorflow_test",
     srcs = ["tensorflow_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",
@@ -171,6 +180,7 @@ go_test(
 go_test(
     name = "wordpress_test",
     srcs = ["wordpress_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",
@@ -190,6 +200,7 @@ go_test(
 go_test(
     name = "pytorch_test",
     srcs = ["pytorch_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",
@@ -210,6 +221,7 @@ go_test(
     embedsrcs = [
         "//test/kubernetes/benchmarks/resources:files",  # keep
     ],
+    nogo = False,
     tags = [
         "local",
         "noguitar",
@@ -230,6 +242,7 @@ go_test(
 go_test(
     name = "stablediffusion_test",
     srcs = ["stablediffusion_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",
@@ -248,6 +261,7 @@ go_test(
 go_test(
     name = "gsutil_test",
     srcs = ["gsutil_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",

test/kubernetes/k8sctx/BUILD

@@ -7,10 +7,12 @@ package(
 
 go_library(
     name = "k8sctx",
+    testonly = True,
     srcs = [
         "k8sctx.go",
         "k8sctx_impl.go",
     ],
+    nogo = False,
     visibility = [
         "//visibility:public",
     ],

test/kubernetes/testcluster/testcluster.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"reflect"
 	"strconv"
 	"strings"
 	"time"
@@ -132,8 +133,31 @@ type TestCluster struct {
 	testNodepoolRuntimeOverride RuntimeType
 }
 
+type testClusterConstructorKey int
+
+const (
+	// testClusterConstructor is the key for the context value that holds the
+	// constructor function for TestCluster.
+	// Defaults to newTestCluster.
+	testClusterConstructor testClusterConstructorKey = iota
+)
+
+// WithTestClusterConstructor returns a context that contains a custom
+// constructor for TestCluster.
+func WithTestClusterConstructor(ctx context.Context, constructor func(context.Context, *testpb.Cluster) (*TestCluster, error)) context.Context {
+	return context.WithValue(ctx, testClusterConstructor, constructor)
+}
+
 // NewTestCluster returns a new TestCluster client.
-func NewTestCluster(cluster *testpb.Cluster) (*TestCluster, error) {
+func NewTestCluster(ctx context.Context, cluster *testpb.Cluster) (*TestCluster, error) {
+	constructor, ok := ctx.Value(testClusterConstructor).(func(context.Context, *testpb.Cluster) (*TestCluster, error))
+	if !ok || constructor == nil || reflect.ValueOf(constructor).IsNil() {
+		constructor = newTestCluster
+	}
+	return constructor(ctx, cluster)
+}
+
+func newTestCluster(_ context.Context, cluster *testpb.Cluster) (*TestCluster, error) {
 	config, err := clientcmd.BuildConfigFromFlags("" /*masterURL*/, cluster.GetCredentialFile())
 	if err != nil {
 		return nil, fmt.Errorf("BuildConfigFromFlags: %w", err)

test/kubernetes/tests/BUILD

@@ -8,6 +8,7 @@ package(
 go_test(
     name = "hello_test",
     srcs = ["hello_test.go"],
+    nogo = False,
     tags = [
         "local",
         "noguitar",