From aee6c515156cfa2675087b2c44069832e9cace97 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Wed, 13 Nov 2024 12:17:31 -0800
Subject: [PATCH] Kubernetes tests: Make `testcluster.NewTestCluster`
 overridable via context.

PiperOrigin-RevId: 696227388
---
 test/kubernetes/testcluster/testcluster.go | 26 +++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/test/kubernetes/testcluster/testcluster.go b/test/kubernetes/testcluster/testcluster.go
index 273b7d1ba..fde5f3f71 100644
--- a/test/kubernetes/testcluster/testcluster.go
+++ b/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)