Test that Cleanup is properly called from BasicRunner

basic_runner_test.go

@@ -21,11 +21,19 @@ func TestBasicRunner_Run(t *testing.T) {
 	r := &BasicRunner{[]Step{stepA, stepB}}
 	r.Run(data)
 
+	// Test run data
 	expected := []string{"a", "b"}
 	results := data["data"].([]string)
 	if !reflect.DeepEqual(results, expected) {
 		t.Errorf("unexpected result: %#v", results)
 	}
+
+	// Test cleanup data
+	expected = []string{"b", "a"}
+	results = data["cleanup"].([]string)
+	if !reflect.DeepEqual(results, expected) {
+		t.Errorf("unexpected result: %#v", results)
+	}
 }
 
 func TestBasicRunner_Run_Halt(t *testing.T) {
@@ -37,9 +45,17 @@ func TestBasicRunner_Run_Halt(t *testing.T) {
 	r := &BasicRunner{[]Step{stepA, stepB, stepC}}
 	r.Run(data)
 
+	// Test run data
 	expected := []string{"a", "b"}
 	results := data["data"].([]string)
 	if !reflect.DeepEqual(results, expected) {
 		t.Errorf("unexpected result: %#v", results)
 	}
+
+	// Test cleanup data
+	expected = []string{"b", "a"}
+	results = data["cleanup"].([]string)
+	if !reflect.DeepEqual(results, expected) {
+		t.Errorf("unexpected result: %#v", results)
+	}
 }

multistep_test.go

@@ -12,13 +12,7 @@ type TestStepAcc struct {
 }
 
 func (s TestStepAcc) Run(state map[string]interface{}) StepAction {
-	if _, ok := state["data"]; !ok {
-		state["data"] = make([]string, 0, 5)
-	}
-
-	data := state["data"].([]string)
-	data = append(data, s.Data)
-	state["data"] = data
+	s.insertData(state, "data")
 
 	if s.Halt {
 		return ActionHalt
@@ -27,5 +21,16 @@ func (s TestStepAcc) Run(state map[string]interface{}) StepAction {
 	return ActionContinue
 }
 
-func (s TestStepAcc) Cleanup(map[string]interface{}) {
+func (s TestStepAcc) Cleanup(state map[string]interface{}) {
+	s.insertData(state, "cleanup")
+}
+
+func (s TestStepAcc) insertData(state map[string]interface{}, key string) {
+	if _, ok := state[key]; !ok {
+		state[key] = make([]string, 0, 5)
+	}
+
+	data := state[key].([]string)
+	data = append(data, s.Data)
+	state[key] = data
 }