Fix runtime admission flaky test due to race condition

Author: jesusha123

plugin/pkg/admission/runtimeclass/BUILD

@@ -14,9 +14,12 @@ go_library(
         "//staging/src/k8s.io/api/node/v1beta1:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/api/equality:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/api/errors:go_default_library",
+        "//staging/src/k8s.io/apimachinery/pkg/apis/meta/v1:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/admission:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/admission/initializer:go_default_library",
         "//staging/src/k8s.io/client-go/informers:go_default_library",
+        "//staging/src/k8s.io/client-go/kubernetes:go_default_library",
+        "//staging/src/k8s.io/client-go/kubernetes/typed/node/v1beta1:go_default_library",
         "//staging/src/k8s.io/client-go/listers/node/v1beta1:go_default_library",
         "//staging/src/k8s.io/component-base/featuregate:go_default_library",
     ],

plugin/pkg/admission/runtimeclass/admission.go

@@ -29,9 +29,12 @@ import (
 	v1beta1 "k8s.io/api/node/v1beta1"
 	apiequality "k8s.io/apimachinery/pkg/api/equality"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apiserver/pkg/admission"
 	genericadmissioninitailizer "k8s.io/apiserver/pkg/admission/initializer"
 	"k8s.io/client-go/informers"
+	"k8s.io/client-go/kubernetes"
+	nodev1beta1client "k8s.io/client-go/kubernetes/typed/node/v1beta1"
 	nodev1beta1listers "k8s.io/client-go/listers/node/v1beta1"
 	"k8s.io/component-base/featuregate"
 	api "k8s.io/kubernetes/pkg/apis/core"
@@ -58,6 +61,7 @@ func Register(plugins *admission.Plugins) {
 type RuntimeClass struct {
 	*admission.Handler
 	runtimeClassLister nodev1beta1listers.RuntimeClassLister
+	runtimeClassClient nodev1beta1client.RuntimeClassInterface
 
 	inspectedFeatures   bool
 	runtimeClassEnabled bool
@@ -68,6 +72,11 @@ var _ admission.MutationInterface = &RuntimeClass{}
 var _ admission.ValidationInterface = &RuntimeClass{}
 
 var _ genericadmissioninitailizer.WantsExternalKubeInformerFactory = &RuntimeClass{}
+var _ genericadmissioninitailizer.WantsExternalKubeClientSet = &RuntimeClass{}
+
+func (r *RuntimeClass) SetExternalKubeClientSet(client kubernetes.Interface) {
+	r.runtimeClassClient = client.NodeV1beta1().RuntimeClasses()
+}
 
 // InspectFeatureGates allows setting bools without taking a dep on a global variable
 func (r *RuntimeClass) InspectFeatureGates(featureGates featuregate.FeatureGate) {
@@ -97,6 +106,9 @@ func (r *RuntimeClass) ValidateInitialization() error {
 	if r.runtimeClassLister == nil {
 		return fmt.Errorf("missing RuntimeClass lister")
 	}
+	if r.runtimeClassClient == nil {
+		return fmt.Errorf("missing RuntimeClass client")
+	}
 	return nil
 }
 
@@ -111,7 +123,7 @@ func (r *RuntimeClass) Admit(ctx context.Context, attributes admission.Attribute
 		return nil
 	}
 
-	pod, runtimeClass, err := r.prepareObjects(attributes)
+	pod, runtimeClass, err := r.prepareObjects(ctx, attributes)
 	if err != nil {
 		return err
 	}
@@ -139,7 +151,7 @@ func (r *RuntimeClass) Validate(ctx context.Context, attributes admission.Attrib
 		return nil
 	}
 
-	pod, runtimeClass, err := r.prepareObjects(attributes)
+	pod, runtimeClass, err := r.prepareObjects(ctx, attributes)
 	if err != nil {
 		return err
 	}
@@ -160,7 +172,7 @@ func NewRuntimeClass() *RuntimeClass {
 }
 
 // prepareObjects returns pod and runtimeClass types from the given admission attributes
-func (r *RuntimeClass) prepareObjects(attributes admission.Attributes) (pod *api.Pod, runtimeClass *v1beta1.RuntimeClass, err error) {
+func (r *RuntimeClass) prepareObjects(ctx context.Context, attributes admission.Attributes) (pod *api.Pod, runtimeClass *v1beta1.RuntimeClass, err error) {
 	pod, ok := attributes.GetObject().(*api.Pod)
 	if !ok {
 		return nil, nil, apierrors.NewBadRequest("Resource was marked with kind Pod but was unable to be converted")
@@ -173,7 +185,11 @@ func (r *RuntimeClass) prepareObjects(attributes admission.Attributes) (pod *api
 	// get RuntimeClass object
 	runtimeClass, err = r.runtimeClassLister.Get(*pod.Spec.RuntimeClassName)
 	if apierrors.IsNotFound(err) {
-		return pod, nil, admission.NewForbidden(attributes, fmt.Errorf("pod rejected: RuntimeClass %q not found", *pod.Spec.RuntimeClassName))
+		// if not found, our informer cache could be lagging, do a live lookup
+		runtimeClass, err = r.runtimeClassClient.Get(ctx, *pod.Spec.RuntimeClassName, metav1.GetOptions{})
+		if apierrors.IsNotFound(err) {
+			return pod, nil, admission.NewForbidden(attributes, fmt.Errorf("pod rejected: RuntimeClass %q not found", *pod.Spec.RuntimeClassName))
+		}
 	}
 
 	// return the pod and runtimeClass.

plugin/pkg/admission/runtimeclass/admission_test.go

@@ -1,3 +1,4 @@
+
 /*
 Copyright 2019 The Kubernetes Authors.
 
@@ -30,6 +31,13 @@ import (
 	"k8s.io/apiserver/pkg/admission"
 	"k8s.io/apiserver/pkg/authentication/user"
 	"k8s.io/kubernetes/pkg/apis/core"
+	"k8s.io/client-go/kubernetes/fake"
+	api "k8s.io/kubernetes/pkg/apis/core"
+	"k8s.io/client-go/informers"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/component-base/featuregate"
+	"k8s.io/kubernetes/pkg/controller"
+	"k8s.io/kubernetes/pkg/features"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -315,6 +323,151 @@ func NewObjectInterfacesForTest() admission.ObjectInterfaces {
 	return admission.NewObjectInterfacesFromScheme(scheme)
 }
 
+func newRuntimeClassForTest(runtimeClassEnabled bool,
+	featureInspection bool,
+	addLister bool,
+	listerObject *v1beta1.RuntimeClass,
+	addClient bool,
+	clientObject *v1beta1.RuntimeClass) *RuntimeClass {
+	runtimeClass := NewRuntimeClass()
+
+	if featureInspection {
+		relevantFeatures := map[featuregate.Feature]featuregate.FeatureSpec{
+			features.RuntimeClass: {Default: runtimeClassEnabled},
+			features.PodOverhead:  {Default: false},
+		}
+		fg := featuregate.NewFeatureGate()
+		fg.Add(relevantFeatures)
+		runtimeClass.InspectFeatureGates(fg)
+	}
+
+	if addLister {
+		informerFactory := informers.NewSharedInformerFactory(nil, controller.NoResyncPeriodFunc())
+		runtimeClass.SetExternalKubeInformerFactory(informerFactory)
+		if listerObject != nil {
+			informerFactory.Node().V1beta1().RuntimeClasses().Informer().GetStore().Add(listerObject)
+		}
+	}
+
+	if addClient {
+		var client kubernetes.Interface
+		if clientObject != nil {
+		    client = fake.NewSimpleClientset(clientObject)
+		} else {
+			client = fake.NewSimpleClientset()
+		}
+		runtimeClass.SetExternalKubeClientSet(client)
+	}
+
+	return runtimeClass
+}
+
+func TestValidateInitialization(t *testing.T) {
+	tests := []struct {
+		name         string
+		expectError  bool
+		runtimeClass *RuntimeClass
+	}{
+		{
+			name:         "runtimeClass disabled, success",
+			expectError:  false,
+			runtimeClass: newRuntimeClassForTest(false, true, true, nil, true, nil),
+		},
+		{
+			name:         "runtimeClass enabled, success",
+			expectError:  false,
+			runtimeClass: newRuntimeClassForTest(true, true, true, nil, true, nil),
+		},
+		{
+			name:         "runtimeClass enabled, no feature inspection",
+			expectError:  true,
+			runtimeClass: newRuntimeClassForTest(true, false, true, nil, true, nil),
+		},
+		{
+			name:         "runtimeClass enabled, no lister",
+			expectError:  true,
+			runtimeClass: newRuntimeClassForTest(true, true, false, nil, true, nil,),
+		},
+		{
+			name:         "runtimeClass enabled, no client",
+			expectError:  true,
+			runtimeClass: newRuntimeClassForTest(true, true, true, nil, false, nil),
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			err := tc.runtimeClass.ValidateInitialization()
+			if tc.expectError {
+				assert.NotEmpty(t, err)
+			} else {
+				assert.Empty(t, err)
+			}
+		})
+	}
+}
+
+func TestAdmit(t *testing.T) {
+	runtimeClassName := "runtimeClassName"
+
+	rc := &v1beta1.RuntimeClass{
+		ObjectMeta: metav1.ObjectMeta{Name: runtimeClassName},
+	}
+
+	pod := api.Pod{
+		ObjectMeta: metav1.ObjectMeta{Name: "podname"},
+		Spec: api.PodSpec{
+			RuntimeClassName: &runtimeClassName,
+		},
+	}
+
+	attributes := admission.NewAttributesRecord(&pod,
+		nil,
+		api.Kind("kind").WithVersion("version"),
+		"",
+		"",
+		api.Resource("pods").WithVersion("version"),
+		"",
+		admission.Create,
+		nil,
+		false,
+		nil)
+
+
+	tests := []struct {
+		name         string
+		expectError  bool
+		runtimeClass *RuntimeClass
+	}{
+		{
+			name:         "runtimeClass found by lister",
+			expectError:  false,
+			runtimeClass: newRuntimeClassForTest(true, true, true, rc, true, nil),
+		},
+		{
+			name:         "runtimeClass found by client",
+			expectError:  false,
+			runtimeClass: newRuntimeClassForTest(true, true, true, nil, true, rc),
+		},
+		{
+			name:         "runtimeClass not found by lister nor client",
+			expectError:  true,
+			runtimeClass: newRuntimeClassForTest(true, true, true, nil, true, nil),
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			err := tc.runtimeClass.Admit(context.TODO(), attributes, nil)
+			if tc.expectError {
+				assert.NotEmpty(t, err)
+			} else {
+				assert.Empty(t, err)
+			}
+		})
+	}
+}
+
 func TestValidate(t *testing.T) {
 	tests := []struct {
 		name         string