CleanupOnFail option for the upgrade Step [#65]

internal/run/upgrade.go

@@ -19,6 +19,7 @@ type Upgrade struct {
 	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)
 	}

internal/run/upgrade_test.go

@@ -101,6 +101,7 @@ func (suite *UpgradeTestSuite) TestPrepareWithUpgradeFlags() {
 		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",