From eeaacc4a07e9fc8de079cf19751c21bf09c0a7a2 Mon Sep 17 00:00:00 2001
From: Morten Lied Johansen <morten.lied.johansen@finn.no>
Date: Thu, 5 Jan 2017 11:45:05 +0100
Subject: [PATCH] Add `--upgrade` option to `init`

When option is given, any existing Tiller will be upgraded to the current version
---
 cmd/helm/init.go                   | 12 +++++++++++-
 cmd/helm/init_test.go              |  3 ++-
 cmd/helm/installer/install.go      | 19 ++++++++++++++++++-
 cmd/helm/installer/install_test.go | 19 +++++++++++++++++++
 4 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/cmd/helm/init.go b/cmd/helm/init.go
index 11f140067..e57fca534 100644
--- a/cmd/helm/init.go
+++ b/cmd/helm/init.go
@@ -65,6 +65,7 @@ type initCmd struct {
 	image      string
 	clientOnly bool
 	canary     bool
+	upgrade    bool
 	namespace  string
 	dryRun     bool
 	out        io.Writer
@@ -94,6 +95,7 @@ func newInitCmd(out io.Writer) *cobra.Command {
 	f := cmd.Flags()
 	f.StringVarP(&i.image, "tiller-image", "i", "", "override tiller image")
 	f.BoolVar(&i.canary, "canary-image", false, "use the canary tiller image")
+	f.BoolVar(&i.upgrade, "upgrade", false, "upgrade if tiller is already installed")
 	f.BoolVarP(&i.clientOnly, "client-only", "c", false, "if set does not install tiller")
 	f.BoolVar(&i.dryRun, "dry-run", false, "do not install local or remote")
 
@@ -130,7 +132,15 @@ func (i *initCmd) run() error {
 			if !kerrors.IsAlreadyExists(err) {
 				return fmt.Errorf("error installing: %s", err)
 			}
-			fmt.Fprintln(i.out, "Warning: Tiller is already installed in the cluster. (Use --client-only to suppress this message.)")
+			if i.upgrade {
+				if err := installer.Upgrade(i.kubeClient, i.namespace, i.image, i.canary); err != nil {
+					return fmt.Errorf("error when upgrading: %s", err)
+				}
+				fmt.Fprintln(i.out, "\nTiller (the helm server side component) has been upgraded to the current version.")
+			} else {
+				fmt.Fprintln(i.out, "Warning: Tiller is already installed in the cluster.\n"+
+					"(Use --client-only to suppress this message, or --upgrade to upgrade Tiller to the current version.)")
+			}
 		} else {
 			fmt.Fprintln(i.out, "\nTiller (the helm server side component) has been installed into your Kubernetes Cluster.")
 		}
diff --git a/cmd/helm/init_test.go b/cmd/helm/init_test.go
index 7ad617570..7a9a85895 100644
--- a/cmd/helm/init_test.go
+++ b/cmd/helm/init_test.go
@@ -89,7 +89,8 @@ func TestInitCmd_exsits(t *testing.T) {
 	if err := cmd.run(); err != nil {
 		t.Errorf("expected error: %v", err)
 	}
-	expected := "Warning: Tiller is already installed in the cluster. (Use --client-only to suppress this message.)"
+	expected := "Warning: Tiller is already installed in the cluster.\n" +
+		"(Use --client-only to suppress this message, or --upgrade to upgrade Tiller to the current version.)"
 	if !strings.Contains(buf.String(), expected) {
 		t.Errorf("expected %q, got %q", expected, buf.String())
 	}
diff --git a/cmd/helm/installer/install.go b/cmd/helm/installer/install.go
index 09f3f6ea5..06856fcdc 100644
--- a/cmd/helm/installer/install.go
+++ b/cmd/helm/installer/install.go
@@ -43,15 +43,32 @@ func Install(client extensionsclient.DeploymentsGetter, namespace, image string,
 	return err
 }
 
+// Upgrade uses kubernetes client to upgrade tiller to current version
+//
+// Returns an error if the command failed.
+func Upgrade(client extensionsclient.DeploymentsGetter, namespace, image string, canary bool) error {
+	obj, err := client.Deployments(namespace).Get("tiller-deploy")
+	if err != nil {
+		return err
+	}
+	obj.Spec.Template.Spec.Containers[0].Image = selectImage(image, canary)
+	_, err = client.Deployments(namespace).Update(obj)
+	return err
+}
+
 // deployment gets the deployment object that installs Tiller.
 func deployment(namespace, image string, canary bool) *extensions.Deployment {
+	return generateDeployment(namespace, selectImage(image, canary))
+}
+
+func selectImage(image string, canary bool) string {
 	switch {
 	case canary:
 		image = defaultImage + ":canary"
 	case image == "":
 		image = fmt.Sprintf("%s:%s", defaultImage, version.Version)
 	}
-	return generateDeployment(namespace, image)
+	return image
 }
 
 // DeploymentManifest gets the manifest (as a string) that describes the Tiller Deployment
diff --git a/cmd/helm/installer/install_test.go b/cmd/helm/installer/install_test.go
index dc9114074..dd22e6304 100644
--- a/cmd/helm/installer/install_test.go
+++ b/cmd/helm/installer/install_test.go
@@ -103,3 +103,22 @@ func TestInstall_canary(t *testing.T) {
 		t.Errorf("unexpected error: %#+v", err)
 	}
 }
+
+func TestUpgrade(t *testing.T) {
+	image := "gcr.io/kubernetes-helm/tiller:v2.0.0"
+
+	fc := fake.NewSimpleClientset(deployment(api.NamespaceDefault, "imageToReplace", false))
+	fc.AddReactor("update", "deployments", func(action testcore.Action) (bool, runtime.Object, error) {
+		obj := action.(testcore.CreateAction).GetObject().(*extensions.Deployment)
+		i := obj.Spec.Template.Spec.Containers[0].Image
+		if i != image {
+			t.Errorf("expected image = '%s', got '%s'", image, i)
+		}
+		return true, obj, nil
+	})
+
+	err := Upgrade(fc.Extensions(), api.NamespaceDefault, image, false)
+	if err != nil {
+		t.Errorf("unexpected error: %#+v", err)
+	}
+}