From 8c6c6fbfa51a3a53b3fb033967927b2b59ec361d Mon Sep 17 00:00:00 2001
From: Erin Call <erin@liffft.com>
Date: Tue, 7 Jan 2020 12:53:55 -0800
Subject: [PATCH] CleanupOnFail option for the upgrade Step [#65]

---
 internal/run/upgrade.go      | 24 ++++++++++++++----------
 internal/run/upgrade_test.go | 26 ++++++++++++++------------
 2 files changed, 28 insertions(+), 22 deletions(-)

diff --git a/internal/run/upgrade.go b/internal/run/upgrade.go
index 3efd95b..c239807 100644
--- a/internal/run/upgrade.go
+++ b/internal/run/upgrade.go
@@ -9,16 +9,17 @@ type Upgrade struct {
 	Chart   string
 	Release string
 
-	ChartVersion string
-	DryRun       bool
-	Wait         bool
-	Values       string
-	StringValues string
-	ValuesFiles  []string
-	ReuseValues  bool
-	Timeout      string
-	Force        bool
-	Atomic       bool
+	ChartVersion  string
+	DryRun        bool
+	Wait          bool
+	Values        string
+	StringValues  string
+	ValuesFiles   []string
+	ReuseValues   bool
+	Timeout       string
+	Force         bool
+	Atomic        bool
+	CleanupOnFail bool
 
 	cmd cmd
 }
@@ -69,6 +70,9 @@ func (u *Upgrade) Prepare(cfg Config) error {
 	if u.Atomic {
 		args = append(args, "--atomic")
 	}
+	if u.CleanupOnFail {
+		args = append(args, "--cleanup-on-fail")
+	}
 	if u.Values != "" {
 		args = append(args, "--set", u.Values)
 	}
diff --git a/internal/run/upgrade_test.go b/internal/run/upgrade_test.go
index f3a1fdf..886fb3b 100644
--- a/internal/run/upgrade_test.go
+++ b/internal/run/upgrade_test.go
@@ -89,18 +89,19 @@ func (suite *UpgradeTestSuite) TestPrepareWithUpgradeFlags() {
 	defer suite.ctrl.Finish()
 
 	u := Upgrade{
-		Chart:        "hot_ac",
-		Release:      "maroon_5_memories",
-		ChartVersion: "radio_edit",
-		DryRun:       true,
-		Wait:         true,
-		Values:       "age=35",
-		StringValues: "height=5ft10in",
-		ValuesFiles:  []string{"/usr/local/stats", "/usr/local/grades"},
-		ReuseValues:  true,
-		Timeout:      "sit_in_the_corner",
-		Force:        true,
-		Atomic:       true,
+		Chart:         "hot_ac",
+		Release:       "maroon_5_memories",
+		ChartVersion:  "radio_edit",
+		DryRun:        true,
+		Wait:          true,
+		Values:        "age=35",
+		StringValues:  "height=5ft10in",
+		ValuesFiles:   []string{"/usr/local/stats", "/usr/local/grades"},
+		ReuseValues:   true,
+		Timeout:       "sit_in_the_corner",
+		Force:         true,
+		Atomic:        true,
+		CleanupOnFail: true,
 	}
 
 	cfg := Config{}
@@ -115,6 +116,7 @@ func (suite *UpgradeTestSuite) TestPrepareWithUpgradeFlags() {
 			"--timeout", "sit_in_the_corner",
 			"--force",
 			"--atomic",
+			"--cleanup-on-fail",
 			"--set", "age=35",
 			"--set-string", "height=5ft10in",
 			"--values", "/usr/local/stats",