From e6ab3cfc543bdef0f893cd25e66c74e5bded696e Mon Sep 17 00:00:00 2001
From: yuwenma <yuwenma@google.com>
Date: Fri, 25 Jan 2019 18:22:55 -0800
Subject: [PATCH] [Mitigate KCM CrashLooping] Add unittests for controllers'
 Init function.

---
 cmd/kube-controller-manager/app/core.go      |  5 ++
 cmd/kube-controller-manager/app/core_test.go | 68 ++++++++++++++------
 2 files changed, 52 insertions(+), 21 deletions(-)

diff --git a/cmd/kube-controller-manager/app/core.go b/cmd/kube-controller-manager/app/core.go
index 5111f77a66..220a790954 100644
--- a/cmd/kube-controller-manager/app/core.go
+++ b/cmd/kube-controller-manager/app/core.go
@@ -36,6 +36,7 @@ import (
 	cacheddiscovery "k8s.io/client-go/discovery/cached"
 	"k8s.io/client-go/dynamic"
 	clientset "k8s.io/client-go/kubernetes"
+	restclient "k8s.io/client-go/rest"
 	csiclientset "k8s.io/csi-api/pkg/client/clientset/versioned"
 	"k8s.io/kubernetes/pkg/controller"
 	cloudcontroller "k8s.io/kubernetes/pkg/controller/cloud"
@@ -331,6 +332,10 @@ func startNamespaceController(ctx ControllerContext) (http.Handler, bool, error)
 	nsKubeconfig.QPS *= 20
 	nsKubeconfig.Burst *= 100
 	namespaceKubeClient := clientset.NewForConfigOrDie(nsKubeconfig)
+	return startModifiedNamespaceController(ctx, namespaceKubeClient, nsKubeconfig)
+}
+
+func startModifiedNamespaceController(ctx ControllerContext, namespaceKubeClient clientset.Interface, nsKubeconfig *restclient.Config) (http.Handler, bool, error) {
 
 	dynamicClient, err := dynamic.NewForConfig(nsKubeconfig)
 	if err != nil {
diff --git a/cmd/kube-controller-manager/app/core_test.go b/cmd/kube-controller-manager/app/core_test.go
index ff2149fdd2..3b410e03ae 100644
--- a/cmd/kube-controller-manager/app/core_test.go
+++ b/cmd/kube-controller-manager/app/core_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package app
 
 import (
+	"net/http"
 	"testing"
 	"time"
 
@@ -34,8 +35,11 @@ type TestClientBuilder struct {
 	clientset clientset.Interface
 }
 
-func (TestClientBuilder) Config(name string) (*restclient.Config, error)  { return nil, nil }
-func (TestClientBuilder) ConfigOrDie(name string) *restclient.Config      { return nil }
+func (TestClientBuilder) Config(name string) (*restclient.Config, error) { return nil, nil }
+func (TestClientBuilder) ConfigOrDie(name string) *restclient.Config {
+	return &restclient.Config{}
+}
+
 func (TestClientBuilder) Client(name string) (clientset.Interface, error) { return nil, nil }
 func (m TestClientBuilder) ClientOrDie(name string) clientset.Interface {
 	return m.clientset
@@ -62,6 +66,10 @@ func (c *FakeClientSet) Discovery() discovery.DiscoveryInterface {
 	return c.DiscoveryObj
 }
 
+func (c *FakeClientSet) GetPossibleResources() []*metav1.APIResourceList {
+	return c.DiscoveryObj.PossibleResources
+}
+
 // Create a fake Clientset with its Discovery method overridden.
 func NewFakeClientset(fakeDiscovery FakeDiscoveryWithError) *FakeClientSet {
 	cs := &FakeClientSet{}
@@ -69,7 +77,30 @@ func NewFakeClientset(fakeDiscovery FakeDiscoveryWithError) *FakeClientSet {
 	return cs
 }
 
-func TestStartResourceQuotaController_DiscoveryError(t *testing.T) {
+func possibleDiscoveryResource() []*metav1.APIResourceList {
+	return []*metav1.APIResourceList{
+		{
+			GroupVersion: "create/v1",
+			APIResources: []metav1.APIResource{
+				{
+					Name:       "jobs",
+					Verbs:      []string{"create", "list", "watch", "delete"},
+					ShortNames: []string{"jz"},
+					Categories: []string{"all"},
+				},
+			},
+		},
+	}
+}
+
+type controllerInitFunc func(ControllerContext) (http.Handler, bool, error)
+
+func TestController_DiscoveryError(t *testing.T) {
+	controllerInitFuncMap := map[string]controllerInitFunc{
+		"ResourceQuotaController":    startResourceQuotaController,
+		"GarbageCollectorController": startGarbageCollectorController,
+	}
+
 	tcs := map[string]struct {
 		discoveryError    error
 		expectedErr       bool
@@ -77,25 +108,13 @@ func TestStartResourceQuotaController_DiscoveryError(t *testing.T) {
 	}{
 		"No Discovery Error": {
 			discoveryError:    nil,
-			possibleResources: nil,
+			possibleResources: possibleDiscoveryResource(),
 			expectedErr:       false,
 		},
 		"Discovery Calls Partially Failed": {
-			discoveryError: new(discovery.ErrGroupDiscoveryFailed),
-			possibleResources: []*metav1.APIResourceList{
-				{
-					GroupVersion: "create/v1",
-					APIResources: []metav1.APIResource{
-						{
-							Name:       "jobs",
-							Verbs:      []string{"create", "list", "watch", "delete"},
-							ShortNames: []string{"jz"},
-							Categories: []string{"all"},
-						},
-					},
-				},
-			},
-			expectedErr: false,
+			discoveryError:    new(discovery.ErrGroupDiscoveryFailed),
+			possibleResources: possibleDiscoveryResource(),
+			expectedErr:       false,
 		},
 	}
 	for name, test := range tcs {
@@ -107,9 +126,16 @@ func TestStartResourceQuotaController_DiscoveryError(t *testing.T) {
 			InformerFactory:  informers.NewSharedInformerFactoryWithOptions(testClientset, time.Duration(1)),
 			InformersStarted: make(chan struct{}),
 		}
-		_, _, err := startResourceQuotaController(ctx)
+		for funcName, controllerInit := range controllerInitFuncMap {
+			_, _, err := controllerInit(ctx)
+			if test.expectedErr != (err != nil) {
+				t.Errorf("%v test failed for use case: %v", funcName, name)
+			}
+		}
+		_, _, err := startModifiedNamespaceController(
+			ctx, testClientset, testClientBuilder.ConfigOrDie("namespace-controller"))
 		if test.expectedErr != (err != nil) {
-			t.Errorf("test failed for use case: %v", name)
+			t.Errorf("Namespace Controller test failed for use case: %v", name)
 		}
 	}
 }