Kubernetes tests & benchmark: Simplify KubernetesContext interface.

This removes the logic around test registration (ended up not needing it),
and simplifies the interface for acquiring/releasing to a single method
that makes the single-cluster implementation more trivial.

PiperOrigin-RevId: 709186717

test/kubernetes/benchmarks/BUILD

@@ -58,6 +58,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )
@@ -88,6 +89,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )
@@ -120,6 +122,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )
@@ -152,6 +155,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )
@@ -184,6 +188,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )
@@ -217,6 +222,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )
@@ -248,6 +254,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )
@@ -280,6 +287,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )
@@ -312,6 +320,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )
@@ -343,6 +352,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )
@@ -382,6 +392,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )
@@ -415,6 +426,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )
@@ -446,6 +458,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )
@@ -478,6 +491,7 @@ go_test(
     ],
     deps = [
         "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/autok8sctx",
         "//test/kubernetes/testcluster",
     ],
 )

test/kubernetes/benchmarks/abslbuild_test.go

@@ -19,26 +19,21 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
 // TestABSLBuild benchmarks building various Abseil C++ targets.
 func TestABSLBuild(t *testing.T) {
 	ctx := context.Background()
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run("ABSL", func(t *testing.T) {
 			t.Parallel()
 			BuildABSL(ctx, t, k8sCtx, cluster)
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestABSLBuild": TestABSLBuild,
-	})
-}

test/kubernetes/benchmarks/ffmpeg_test.go

@@ -19,25 +19,20 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
 func TestFfmpeg(t *testing.T) {
 	ctx := context.Background()
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run("ffmpeg", func(t *testing.T) {
 			t.Parallel()
 			RunFFMPEG(ctx, t, k8sCtx, cluster)
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestFfmpeg": TestFfmpeg,
-	})
-}

test/kubernetes/benchmarks/grpc_test.go

@@ -19,25 +19,20 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
 func TestGRPCBuild(t *testing.T) {
 	ctx := context.Background()
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run("gRPC", func(t *testing.T) {
 			t.Parallel()
 			BuildGRPC(ctx, t, k8sCtx, cluster)
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestGRPCBuild": TestGRPCBuild,
-	})
-}

test/kubernetes/benchmarks/gsutil_test.go

@@ -21,25 +21,20 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
 func TestGSUtil(t *testing.T) {
 	ctx := context.Background()
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run("GSUtil", func(t *testing.T) {
 			t.Parallel()
 			RunGSUtil(ctx, t, k8sCtx, cluster)
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestGSUtil": TestGSUtil,
-	})
-}

test/kubernetes/benchmarks/nginx_test.go

@@ -19,25 +19,20 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
 func TestNginx(t *testing.T) {
 	ctx := context.Background()
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run("nginx", func(t *testing.T) {
 			t.Parallel()
 			BenchmarkNginx(ctx, t, k8sCtx, cluster)
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestNginx": TestNginx,
-	})
-}

test/kubernetes/benchmarks/ollama_test.go

@@ -22,6 +22,7 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
@@ -29,20 +30,14 @@ func TestOllama(t *testing.T) {
 	fmt.Fprint(os.Stderr, "HEADS UP: This test uses a huge container image which may take up to 30 minutes to download onto nodes the first time you run it.\n")
 
 	ctx := context.Background()
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run("Ollama", func(t *testing.T) {
 			t.Parallel()
 			BenchmarkOllama(ctx, t, k8sCtx, cluster)
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestOllama": TestOllama,
-	})
-}

test/kubernetes/benchmarks/postgresql_test.go

@@ -20,26 +20,21 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
 // TestPostgresPGBench benchmarks a PostgreSQL database with pgbench.
 func TestPostgresPGBench(t *testing.T) {
 	ctx := context.Background()
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run("PostgresPGBench", func(t *testing.T) {
 			t.Parallel()
 			BenchmarkPostgresPGBench(ctx, t, k8sCtx, cluster)
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestPostgresPGBench": TestPostgresPGBench,
-	})
-}

test/kubernetes/benchmarks/pytorch_test.go

@@ -19,6 +19,7 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
@@ -48,24 +49,14 @@ func TestMobileNetV2(t *testing.T) {
 }
 
 func runTests(ctx context.Context, t *testing.T, tests []pytorchTest) {
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run("PyTorch", func(t *testing.T) {
 			t.Parallel()
 			RunPytorch(ctx, t, k8sCtx, cluster, tests)
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestFastNLPBert":       TestFastNLPBert,
-		"TestBigBird":           TestBigBird,
-		"TestSpeechTransformer": TestSpeechTransformer,
-		"TestLearningToPaint":   TestLearningToPaint,
-		"TestMobileNetV2":       TestMobileNetV2,
-	})
-}

test/kubernetes/benchmarks/redis_test.go

@@ -19,26 +19,21 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
 // TestRedis benchmarks redis servers on k8s clusters.
 func TestRedis(t *testing.T) {
 	ctx := context.Background()
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run("Redis", func(t *testing.T) {
 			t.Parallel()
 			BenchmarkRedis(ctx, t, k8sCtx, cluster)
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestRedis": TestRedis,
-	})
-}

test/kubernetes/benchmarks/rubydev_test.go

@@ -19,26 +19,21 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
 // TestRubyDev benchmarks a build job on k8s clusters.
 func TestRubyDev(t *testing.T) {
 	ctx := context.Background()
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run("RubyDev", func(t *testing.T) {
 			t.Parallel()
 			RunRubyDev(ctx, t, k8sCtx, cluster)
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestRubyDev": TestRubyDev,
-	})
-}

test/kubernetes/benchmarks/stablediffusion_test.go

@@ -19,25 +19,20 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
 func TestStableDiffusionXL(t *testing.T) {
 	ctx := context.Background()
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run("stable_diffusion_xl", func(t *testing.T) {
 			t.Parallel()
 			RunStableDiffusionXL(ctx, t, k8sCtx, cluster)
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestStableDiffusionXL": TestStableDiffusionXL,
-	})
-}

test/kubernetes/benchmarks/startup_test.go

@@ -19,16 +19,17 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
 func TestStartup(t *testing.T) {
 	ctx := context.Background()
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run(benchName, func(t *testing.T) {
 			cluster := cluster
 			t.Parallel()
@@ -36,9 +37,3 @@ func TestStartup(t *testing.T) {
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestStartup": TestStartup,
-	})
-}

test/kubernetes/benchmarks/tensorflow_test.go

@@ -19,25 +19,20 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
 func TestTensorflowOnCPU(t *testing.T) {
 	ctx := context.Background()
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run("TensorflowOnCPU", func(t *testing.T) {
 			t.Parallel()
 			RunTensorflowOnCPU(ctx, t, k8sCtx, cluster)
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestTensorflowOnCPU": TestTensorflowOnCPU,
-	})
-}

test/kubernetes/benchmarks/wordpress_test.go

@@ -19,25 +19,20 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx/autok8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 )
 
 func TestWordpress(t *testing.T) {
 	ctx := context.Background()
-	k8sCtx, err := k8sctx.Context(ctx)
+	k8sCtx, err := autok8sctx.New(ctx)
 	if err != nil {
 		t.Fatalf("Failed to get kubernetes context: %v", err)
 	}
-	k8sCtx.ForEachCluster(ctx, t, func(cluster *testcluster.TestCluster) {
+	k8sctx.ForEachCluster(ctx, t, k8sCtx, func(cluster *testcluster.TestCluster) {
 		t.Run("wordpress", func(t *testing.T) {
 			t.Parallel()
 			BenchmarkWordpress(ctx, t, k8sCtx, cluster)
 		})
 	})
 }
-
-func TestMain(m *testing.M) {
-	k8sctx.TestMain(m, map[string]k8sctx.TestFunc{
-		"TestWordpress": TestWordpress,
-	})
-}

test/kubernetes/k8sctx/BUILD

@@ -10,17 +10,12 @@ go_library(
     testonly = True,
     srcs = [
         "k8sctx.go",
-        "k8sctx_impl.go",
     ],
     nogo = False,
     visibility = [
         "//visibility:public",
     ],
     deps = [
-        "//runsc/flag",
         "//test/kubernetes/testcluster",
-        "//tools/gvisor_k8s_tool/provider/kubectl",
-        "@org_golang_google_protobuf//encoding/prototext:go_default_library",
-        "@org_golang_google_protobuf//types/known/anypb:go_default_library",
     ],
 )

test/kubernetes/k8sctx/autok8sctx/BUILD

@@ -0,0 +1,22 @@
+load("//tools:defs.bzl", "go_library")
+
+package(
+    default_applicable_licenses = ["//:license"],
+    licenses = ["notice"],
+)
+
+go_library(
+    name = "autok8sctx",
+    testonly = True,
+    srcs = ["kubectlctx.go"],
+    nogo = False,
+    visibility = [
+        "//visibility:public",
+    ],
+    deps = [
+        "//runsc/flag",
+        "//test/kubernetes/k8sctx",
+        "//test/kubernetes/testcluster",
+        "//tools/gvisor_k8s_tool/provider/kubectl",
+    ],
+)

test/kubernetes/k8sctx/autok8sctx/kubectlctx.go

@@ -0,0 +1,50 @@
+// Copyright 2024 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build !false
+// +build !false
+
+// Package autok8sctx provides a KubernetesContext that uses a kubectl config
+// and context to determine the test cluster to use for tests and benchmarks.
+// See parent package (`k8sctx`) for more info.
+package autok8sctx
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"gvisor.dev/gvisor/runsc/flag"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
+	"gvisor.dev/gvisor/test/kubernetes/testcluster"
+	"gvisor.dev/gvisor/tools/gvisor_k8s_tool/provider/kubectl"
+)
+
+var (
+	kubectlContextName = flag.String("kubectl-context-name", "", "Name of the kubectl context to use within the kubectl config")
+)
+
+// New creates a KubernetesContext using flags to determine which kubectl
+// config and kubectl context to use as the test cluster.
+func New(ctx context.Context) (k8sctx.KubernetesContext, error) {
+	if *kubectlContextName == "" {
+		return nil, errors.New("no kubectl context name specified")
+	}
+	cluster, err := kubectl.NewCluster(*kubectlContextName)
+	if err != nil {
+		return nil, fmt.Errorf("cannot initialize cluster %q: %w", *kubectlContextName, err)
+	}
+	testCluster := testcluster.NewTestClusterFromClient(*kubectlContextName, cluster.Client())
+	return k8sctx.NewSingleCluster(testCluster), nil
+}

test/kubernetes/k8sctx/k8sctx.go

@@ -21,8 +21,6 @@ package k8sctx
 
 import (
 	"context"
-	"errors"
-	"fmt"
 	"sync"
 	"testing"
 
@@ -34,29 +32,11 @@ import (
 // Tests are expected to call `RegisterTest` for every of their test function,
 // then `TestMain`.
 type KubernetesContext interface {
-	// TestMain should be called inside tests' `TestMain` function, after having
-	// registered all tests with `RegisterTest`.
-	TestMain(m *testing.M)
-
-	// RegisterTest registers a test.
-	// It should be called for every `Test*(*testing.T)` function in the test.
-	// Note that the `k8sctx.TestMain` helper function below will call this for
-	// you given a map of tests.
-	RegisterTest(name string, fn TestFunc)
-
-	// AcquireCluster returns a single cluster for the test or benchmark to use.
+	// Cluster returns a single cluster for the test or benchmark to use.
 	// The cluster is guaranteed to not be in use by other tests or benchmarks
-	// until the `ReleaseCluster` method is called.
-	// This method should block if there are no available clusters.
-	AcquireCluster(ctx context.Context, t *testing.T) *testcluster.TestCluster
-
-	// ReleaseCluster unlocks the given cluster for use by other tests or
-	// benchmarks.
-	ReleaseCluster(ctx context.Context, t *testing.T, cluster *testcluster.TestCluster)
-
-	// ForEachCluster reserves as many test clusters as are available, calls
-	// `fn` on each of them, and releases each of them when `fn` finishes.
-	ForEachCluster(ctx context.Context, t *testing.T, fn func(cluster *testcluster.TestCluster))
+	// until the returned function is called.
+	// If there are no available clusters, it returns a nil TestCluster.
+	Cluster(ctx context.Context, t *testing.T) (*testcluster.TestCluster, func())
 
 	// ResolveImage resolves a container image name (possibly with a label)
 	// to a fully-qualified image name. It can also return an `image:label`
@@ -65,47 +45,55 @@ type KubernetesContext interface {
 	ResolveImage(ctx context.Context, imageName string) (string, error)
 }
 
-// TestFunc is a test function that is expected to call `Context` and run a
-// test or benchmark within a Kubernetes context.
-type TestFunc func(t *testing.T)
-
-var (
-	kubernetesCtxMu   sync.Mutex
-	kubernetesCtxOnce sync.Once
-	kubernetesCtxFn   func(context.Context) (KubernetesContext, error)
-	kubernetesCtx     KubernetesContext
-	kubernetesCtxErr  error
-)
-
-// Context gets the global Kubernetes context.
-// It must be called after SetContext has already been called.
-func Context(ctx context.Context) (KubernetesContext, error) {
-	kubernetesCtxMu.Lock()
-	defer kubernetesCtxMu.Unlock()
-	if kubernetesCtxFn == nil {
-		return nil, errors.New("k8sctx.Context called prior to k8sctx.SetContextConstructor")
+// ForEachCluster calls the given function for each available cluster
+// sequentially.
+// In order to run per-cluster subtests in parallel, call `t.Run` inside
+// `fn` and then `t.Parallel` inside that.
+func ForEachCluster(ctx context.Context, t *testing.T, k8sCtx KubernetesContext, fn func(cluster *testcluster.TestCluster)) {
+	var clusterFns []func()
+	for {
+		cluster, releaseFn := k8sCtx.Cluster(ctx, t)
+		if cluster == nil {
+			break
+		}
+		clusterFns = append(clusterFns, func() {
+			defer releaseFn()
+			fn(cluster)
+		})
+	}
+	for _, clusterFn := range clusterFns {
+		clusterFn()
 	}
-	kubernetesCtxOnce.Do(func() {
-		kubernetesCtx, kubernetesCtxErr = kubernetesCtxFn(ctx)
-	})
-	return kubernetesCtx, kubernetesCtxErr
 }
 
-// SetContextConstructor sets the global Kubernetes context constructor.
-func SetContextConstructor(fn func(context.Context) (KubernetesContext, error)) {
-	kubernetesCtxMu.Lock()
-	defer kubernetesCtxMu.Unlock()
-	kubernetesCtxFn = fn
+// singleCluster implements KubernetesContext using a single cluster.
+type singleCluster struct {
+	mu      sync.Mutex
+	cluster *testcluster.TestCluster
 }
 
-// TestMain is a helper to write the TestMain function of tests.
-func TestMain(m *testing.M, testFuncs map[string]TestFunc) {
-	k8sCtx, err := Context(context.Background())
-	if err != nil {
-		panic(fmt.Sprintf("failed to get k8sctx: %v", err))
+// Cluster implements KubernetesContext.Cluster.
+func (sc *singleCluster) Cluster(ctx context.Context, t *testing.T) (*testcluster.TestCluster, func()) {
+	sc.mu.Lock()
+	defer sc.mu.Unlock()
+	cl := sc.cluster
+	sc.cluster = nil
+	return cl, func() {
+		if cl != nil {
+			sc.mu.Lock()
+			defer sc.mu.Unlock()
+			sc.cluster = cl
+		}
 	}
-	for name, fn := range testFuncs {
-		k8sCtx.RegisterTest(name, fn)
-	}
-	k8sCtx.TestMain(m)
+}
+
+// ResolveImage implements KubernetesContext.ResolveImage.
+func (*singleCluster) ResolveImage(ctx context.Context, imageName string) (string, error) {
+	return imageName, nil
+}
+
+// NewSingleCluster creates a KubernetesContext that uses a single, static
+// test cluster.
+func NewSingleCluster(cluster *testcluster.TestCluster) KubernetesContext {
+	return &singleCluster{cluster: cluster}
 }

test/kubernetes/k8sctx/k8sctx_impl.go

@@ -1,98 +0,0 @@
-// Copyright 2024 The gVisor Authors.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-//go:build !false
-// +build !false
-
-package k8sctx
-
-import (
-	"context"
-	"errors"
-	"fmt"
-	"os"
-	"testing"
-
-	"google.golang.org/protobuf/encoding/prototext"
-	"google.golang.org/protobuf/types/known/anypb"
-	"gvisor.dev/gvisor/runsc/flag"
-	"gvisor.dev/gvisor/test/kubernetes/testcluster"
-	"gvisor.dev/gvisor/tools/gvisor_k8s_tool/provider/kubectl"
-)
-
-var (
-	kubectlContextName  = flag.String("kubectl-context-name", "", "Name of the kubectl context to use within the kubectl config")
-	clusterProtoPath    = flag.String("cluter-proto-path", "", "Path to a `google.container.v1.Cluster` textproto file")
-	testNodepoolRuntime = flag.String("test-nodepool-runtime", "", "if set, override the runtime used for pods scheduled on the 'test' nodepool. If unset, the nodepool default is used")
-)
-
-// kubectlContext implements KubernetesContext using a named `kubectl` context
-// from the user's kubectl config.
-type kubectlContext struct {
-	cluster *testcluster.TestCluster
-}
-
-func newKubectlContext(ctx context.Context) (KubernetesContext, error) {
-	if *kubectlContextName == "" {
-		return nil, errors.New("no kubectl context name specified")
-	}
-	if *clusterProtoPath == "" {
-		return nil, errors.New("no cluster proto path specified")
-	}
-	cluster, err := kubectl.NewCluster(*kubectlContextName)
-	if err != nil {
-		return nil, fmt.Errorf("cannot initialize cluster %q: %w", *kubectlContextName, err)
-	}
-	var clusterPB anypb.Any
-	clusterBytes, err := os.ReadFile(*clusterProtoPath)
-	if err != nil {
-		return nil, fmt.Errorf("cannot read cluster textproto file %q: %w", *clusterProtoPath, err)
-	}
-	if err = prototext.Unmarshal(clusterBytes, &clusterPB); err != nil {
-		return nil, fmt.Errorf("cannot unmarshal cluster textproto file %q: %w", *clusterProtoPath, err)
-	}
-	testCluster := testcluster.NewTestClusterFromClient(*kubectlContextName, cluster.Client())
-	if *testNodepoolRuntime != "" {
-		testCluster.OverrideTestNodepoolRuntime(testcluster.RuntimeType(*testNodepoolRuntime))
-	}
-	return &kubectlContext{cluster: testCluster}, nil
-}
-
-func (c *kubectlContext) AcquireCluster(ctx context.Context, t *testing.T) *testcluster.TestCluster {
-	return c.cluster
-}
-
-func (c *kubectlContext) ReleaseCluster(ctx context.Context, t *testing.T, cluster *testcluster.TestCluster) {
-	// Nothing to do.
-}
-
-func (c *kubectlContext) ForEachCluster(ctx context.Context, t *testing.T, fn func(cluster *testcluster.TestCluster)) {
-	fn(c.cluster)
-}
-
-func (c *kubectlContext) ResolveImage(ctx context.Context, imageName string) (string, error) {
-	return imageName, nil
-}
-
-func (c *kubectlContext) RegisterTest(name string, fn TestFunc) {
-	// Nothing to do here, we use the regular testing library.
-}
-
-func (c *kubectlContext) TestMain(m *testing.M) {
-	os.Exit(m.Run())
-}
-
-func init() {
-	SetContextConstructor(newKubectlContext)
-}