Updating test infra to work for multiple tests

Author: kfswain

pkg/epp/server/controller_manager.go

@@ -28,6 +28,7 @@ import (
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/cache"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/manager"
 	"sigs.k8s.io/gateway-api-inference-extension/api/v1alpha2"
 )
 
@@ -40,7 +41,7 @@ func init() {
 
 // NewDefaultManager creates a new controller manager with default configuration.
 func NewDefaultManager(namespace, name string, restConfig *rest.Config) (ctrl.Manager, error) {
-	manager, err := ctrl.NewManager(restConfig, ctrl.Options{
+	defaultOpts := ctrl.Options{
 		Scheme: scheme,
 		Cache: cache.Options{
 			ByObject: map[client.Object]cache.ByObject{
@@ -65,7 +66,13 @@ func NewDefaultManager(namespace, name string, restConfig *rest.Config) (ctrl.Ma
 				},
 			},
 		},
-	})
+	}
+	return NewManagerWithOptions(restConfig, defaultOpts)
+}
+
+// NewManagerWithOptions creates a new controller manager with injectable options.
+func NewManagerWithOptions(restConfig *rest.Config, opts manager.Options) (ctrl.Manager, error) {
+	manager, err := ctrl.NewManager(restConfig, opts)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create controller manager: %v", err)
 	}

test/integration/epp/hermetic_test.go

@@ -43,6 +43,8 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/protobuf/testing/protocmp"
 	"google.golang.org/protobuf/types/known/structpb"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/fields"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
@@ -51,7 +53,10 @@ import (
 	"k8s.io/component-base/metrics/legacyregistry"
 	metricsutils "k8s.io/component-base/metrics/testutil"
 	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/cache"
+	"sigs.k8s.io/controller-runtime/pkg/client"
 	k8sclient "sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/config"
 	"sigs.k8s.io/controller-runtime/pkg/envtest"
 	"sigs.k8s.io/controller-runtime/pkg/manager"
 	"sigs.k8s.io/gateway-api-inference-extension/api/v1alpha2"
@@ -78,6 +83,13 @@ var (
 	logger       = logutil.NewTestLogger().V(logutil.VERBOSE)
 )
 
+func TestMain(m *testing.M) {
+	cleanup := BeforeSuite()
+	code := m.Run()
+	cleanup()
+	os.Exit(code)
+}
+
 func TestKubeInferenceModelRequest(t *testing.T) {
 	tests := []struct {
 		name              string
@@ -337,11 +349,6 @@ func TestKubeInferenceModelRequest(t *testing.T) {
 			wantErr: false,
 		},
 	}
-
-	// Set up global k8sclient and extproc server runner with test environment config
-	cleanup := BeforeSuit(t)
-	defer cleanup()
-
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			client, cleanup := setUpHermeticServer(t, test.pods, false)
@@ -1266,10 +1273,6 @@ func TestFullDuplexStreamed_KubeInferenceModelRequest(t *testing.T) {
 		},
 	}
 
-	// Set up global k8sclient and extproc server runner with test environment config
-	cleanup := BeforeSuit(t)
-	defer cleanup()
-
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			client, cleanup := setUpHermeticServer(t, test.pods, true)
@@ -1380,7 +1383,7 @@ func fakePod(index int) backendmetrics.Pod {
 }
 
 // Sets up a test environment and returns the runner struct
-func BeforeSuit(t *testing.T) func() {
+func BeforeSuite() func() {
 	// Set up mock k8s API Client
 	testEnv = &envtest.Environment{
 		CRDDirectoryPaths:     []string{filepath.Join("..", "..", "..", "config", "crd", "bases")},
@@ -1404,7 +1407,7 @@ func BeforeSuit(t *testing.T) func() {
 	// Init runtime.
 	ctrl.SetLogger(logger)
 
-	mgr, err := server.NewDefaultManager("default", "vllm-llama2-7b-pool", cfg)
+	mgr, err := server.NewManagerWithOptions(cfg, managerTestOptions("default", "vllm-llama2-7b-pool"))
 	if err != nil {
 		logutil.Fatal(logger, err, "Failed to create controller manager")
 	}
@@ -1425,7 +1428,7 @@ func BeforeSuit(t *testing.T) func() {
 		logutil.Fatal(logger, err, "Failed to setup server runner")
 	}
 
-	// Start the controller manager in go routine, not blocking
+	// Start the controller manager in a go routine, not blocking
 	go func() {
 		if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
 			logutil.Fatal(logger, err, "Failed to start manager")
@@ -1466,14 +1469,28 @@ func BeforeSuit(t *testing.T) func() {
 		}
 	}
 
-	assert.EventuallyWithT(t, func(t *assert.CollectT) {
+	var simpleTimeout bool
+	go func() {
+		time.Sleep(10 * time.Second)
+		simpleTimeout = true
+	}()
+
+	for {
+		if simpleTimeout {
+			logutil.Fatal(logger, nil, "Sync failed")
+		}
 		modelExist := serverRunner.Datastore.ModelGet("my-model")
 		synced := serverRunner.Datastore.PoolHasSynced() && modelExist != nil
-		assert.True(t, synced, "Timeout waiting for the pool and models to sync")
-	}, 10*time.Second, 10*time.Millisecond)
+		if synced {
+			break
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
 
 	return func() {
 		_ = testEnv.Stop()
+		_ = k8sClient.DeleteAllOf(context.Background(), &v1alpha2.InferencePool{})
+		_ = k8sClient.DeleteAllOf(context.Background(), &v1alpha2.InferenceModel{})
 	}
 }
 
@@ -1601,3 +1618,41 @@ func registerMetricsHandler(mgr manager.Manager, port int) error {
 	}
 	return nil
 }
+
+// inject options that allow multiple test runs to run
+// https://github.com/kubernetes-sigs/controller-runtime/issues/2937
+func managerTestOptions(namespace, name string) ctrl.Options {
+	return ctrl.Options{
+		Scheme: scheme,
+		Cache: cache.Options{
+			ByObject: map[client.Object]cache.ByObject{
+				&corev1.Pod{}: {
+					Namespaces: map[string]cache.Config{
+						namespace: {},
+					},
+				},
+				&v1alpha2.InferencePool{}: {
+					Namespaces: map[string]cache.Config{
+						namespace: {
+							FieldSelector: fields.SelectorFromSet(fields.Set{
+								"metadata.name": name,
+							}),
+						},
+					},
+				},
+				&v1alpha2.InferenceModel{}: {
+					Namespaces: map[string]cache.Config{
+						namespace: {},
+					},
+				},
+			},
+		},
+		Controller: config.Controller{
+			SkipNameValidation: boolPointer(true),
+		},
+	}
+}
+
+func boolPointer(b bool) *bool {
+	return &b
+}