Replace type switches in Rollback with Visitor pattern

pkg/kubectl/BUILD

@@ -23,6 +23,7 @@ go_test(
         "quota_test.go",
         "resource_filter_test.go",
         "rolebinding_test.go",
+        "rollback_test.go",
         "rolling_updater_test.go",
         "rollout_status_test.go",
         "run_test.go",

pkg/kubectl/rollback.go

@@ -35,13 +35,13 @@ import (
 	"k8s.io/apimachinery/pkg/watch"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/kubernetes/pkg/api/legacyscheme"
-	"k8s.io/kubernetes/pkg/apis/apps"
 	api "k8s.io/kubernetes/pkg/apis/core"
 	apiv1 "k8s.io/kubernetes/pkg/apis/core/v1"
 	"k8s.io/kubernetes/pkg/apis/extensions"
 	"k8s.io/kubernetes/pkg/controller/daemon"
 	deploymentutil "k8s.io/kubernetes/pkg/controller/deployment/util"
 	"k8s.io/kubernetes/pkg/controller/statefulset"
+	kapps "k8s.io/kubernetes/pkg/kubectl/apps"
 	sliceutil "k8s.io/kubernetes/pkg/kubectl/util/slice"
 	printersinternal "k8s.io/kubernetes/pkg/printers/internalversion"
 )
@@ -56,16 +56,46 @@ type Rollbacker interface {
 	Rollback(obj runtime.Object, updatedAnnotations map[string]string, toRevision int64, dryRun bool) (string, error)
 }
 
+type RollbackVisitor struct {
+	clientset kubernetes.Interface
+	result    Rollbacker
+}
+
+func (v *RollbackVisitor) VisitDeployment(elem kapps.GroupKindElement) {
+	v.result = &DeploymentRollbacker{v.clientset}
+}
+
+func (v *RollbackVisitor) VisitStatefulSet(kind kapps.GroupKindElement) {
+	v.result = &StatefulSetRollbacker{v.clientset}
+}
+
+func (v *RollbackVisitor) VisitDaemonSet(kind kapps.GroupKindElement) {
+	v.result = &DaemonSetRollbacker{v.clientset}
+}
+
+func (v *RollbackVisitor) VisitJob(kind kapps.GroupKindElement)                   {}
+func (v *RollbackVisitor) VisitPod(kind kapps.GroupKindElement)                   {}
+func (v *RollbackVisitor) VisitReplicaSet(kind kapps.GroupKindElement)            {}
+func (v *RollbackVisitor) VisitReplicationController(kind kapps.GroupKindElement) {}
+
+// RollbackerFor returns an implementation of Rollbacker interface for the given schema kind
 func RollbackerFor(kind schema.GroupKind, c kubernetes.Interface) (Rollbacker, error) {
-	switch kind {
-	case extensions.Kind("Deployment"), apps.Kind("Deployment"):
-		return &DeploymentRollbacker{c}, nil
-	case extensions.Kind("DaemonSet"), apps.Kind("DaemonSet"):
-		return &DaemonSetRollbacker{c}, nil
-	case apps.Kind("StatefulSet"):
-		return &StatefulSetRollbacker{c}, nil
+	elem := kapps.GroupKindElement(kind)
+	visitor := &RollbackVisitor{
+		clientset: c,
 	}
-	return nil, fmt.Errorf("no rollbacker has been implemented for %q", kind)
+
+	err := elem.Accept(visitor)
+
+	if err != nil {
+		return nil, fmt.Errorf("error retrieving rollbacker for %q, %v", kind.String(), err)
+	}
+
+	if visitor.result == nil {
+		return nil, fmt.Errorf("no rollbacker has been implemented for %q", kind.String())
+	}
+
+	return visitor.result, nil
 }
 
 type DeploymentRollbacker struct {

pkg/kubectl/rollback_test.go

@@ -0,0 +1,46 @@
+/*
+Copyright 2017 The Kubernetes 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.
+*/
+
+package kubectl
+
+import (
+	"reflect"
+	"testing"
+
+	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/client-go/kubernetes/fake"
+)
+
+var rollbacktests = map[schema.GroupKind]reflect.Type{
+	{Group: "apps", Kind: "DaemonSet"}:   reflect.TypeOf(&DaemonSetRollbacker{}),
+	{Group: "apps", Kind: "StatefulSet"}: reflect.TypeOf(&StatefulSetRollbacker{}),
+	{Group: "apps", Kind: "Deployment"}:  reflect.TypeOf(&DeploymentRollbacker{}),
+}
+
+func TestRollbackerFor(t *testing.T) {
+	fakeClientset := &fake.Clientset{}
+
+	for kind, expectedType := range rollbacktests {
+		result, err := RollbackerFor(kind, fakeClientset)
+		if err != nil {
+			t.Fatalf("error getting Rollbacker for a %v: %v", kind.String(), err)
+		}
+
+		if reflect.TypeOf(result) != expectedType {
+			t.Fatalf("unexpected output type (%v was expected but got %v)", expectedType, reflect.TypeOf(result))
+		}
+	}
+}