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

@@ -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 {

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 {