From d031651c454082afdd05bf00c87123c8a7456007 Mon Sep 17 00:00:00 2001
From: Taylor Thomas <taylor.thomas@microsoft.com>
Date: Wed, 3 Jul 2019 13:05:47 -0600
Subject: [PATCH] fix(pkg/action): Adds back in missing wait functionality

In my previous PR, I did not notice that the wait functionality had been completely removed
from the actions. This restores wait functionality to upgrade and rollback

Signed-off-by: Taylor Thomas <taylor.thomas@microsoft.com>
---
 pkg/action/rollback.go | 12 +++++++++++-
 pkg/action/upgrade.go  | 10 ++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/pkg/action/rollback.go b/pkg/action/rollback.go
index 68f508ceb..2db6ed7a9 100644
--- a/pkg/action/rollback.go
+++ b/pkg/action/rollback.go
@@ -149,7 +149,7 @@ func (r *Rollback) performRollback(currentRelease, targetRelease *release.Releas
 
 	cr := bytes.NewBufferString(currentRelease.Manifest)
 	tr := bytes.NewBufferString(targetRelease.Manifest)
-	// TODO add wait
+
 	if err := r.cfg.KubeClient.Update(cr, tr, r.Force, r.Recreate); err != nil {
 		msg := fmt.Sprintf("Rollback %q failed: %s", targetRelease.Name, err)
 		r.cfg.Log("warning: %s", msg)
@@ -161,6 +161,16 @@ func (r *Rollback) performRollback(currentRelease, targetRelease *release.Releas
 		return targetRelease, err
 	}
 
+	if r.Wait {
+		buf := bytes.NewBufferString(targetRelease.Manifest)
+		if err := r.cfg.KubeClient.Wait(buf, r.Timeout); err != nil {
+			targetRelease.SetStatus(release.StatusFailed, fmt.Sprintf("Release %q failed: %s", targetRelease.Name, err.Error()))
+			r.cfg.recordRelease(currentRelease)
+			r.cfg.recordRelease(targetRelease)
+			return targetRelease, errors.Wrapf(err, "release %s failed", targetRelease.Name)
+		}
+	}
+
 	// post-rollback hooks
 	if !r.DisableHooks {
 		if err := r.execHook(targetRelease.Hooks, hooks.PostRollback); err != nil {
diff --git a/pkg/action/upgrade.go b/pkg/action/upgrade.go
index ecf9883e0..6c2c67a02 100644
--- a/pkg/action/upgrade.go
+++ b/pkg/action/upgrade.go
@@ -211,6 +211,16 @@ func (u *Upgrade) performUpgrade(originalRelease, upgradedRelease *release.Relea
 		return upgradedRelease, err
 	}
 
+	if u.Wait {
+		buf := bytes.NewBufferString(upgradedRelease.Manifest)
+		if err := u.cfg.KubeClient.Wait(buf, u.Timeout); err != nil {
+			upgradedRelease.SetStatus(release.StatusFailed, fmt.Sprintf("Release %q failed: %s", upgradedRelease.Name, err.Error()))
+			u.cfg.recordRelease(originalRelease)
+			u.cfg.recordRelease(upgradedRelease)
+			return upgradedRelease, errors.Wrapf(err, "release %s failed", upgradedRelease.Name)
+		}
+	}
+
 	// post-upgrade hooks
 	if !u.DisableHooks {
 		if err := u.execHook(upgradedRelease.Hooks, hooks.PostUpgrade); err != nil {