Have NewAskStubber perform auto-cleanup

mislav · mislav · commit a33b5a55c437 · 2022-01-14T18:52:00.000+01:00
diff --git a/pkg/cmd/gist/view/view_test.go b/pkg/cmd/gist/view/view_test.go
@@ -355,8 +355,7 @@ func Test_viewRun(t *testing.T) {
 				)),
 			)
 
-			as, surveyteardown := prompt.NewAskStubber()
-			defer surveyteardown(t)
+			as := prompt.NewAskStubber(t)
 			as.StubPrompt("Select a gist").AnswerDefault()
 		}
 
@@ -470,8 +469,7 @@ func Test_promptGists(t *testing.T) {
 		client := &http.Client{Transport: reg}
 
 		t.Run(tt.name, func(t *testing.T) {
-			as, surveyteardown := prompt.NewAskStubber()
-			defer surveyteardown(t)
+			as := prompt.NewAskStubber(t)
 			if tt.askStubs != nil {
 				tt.askStubs(as)
 			}
diff --git a/pkg/cmd/pr/review/review_test.go b/pkg/cmd/pr/review/review_test.go
@@ -309,8 +309,7 @@ func TestPRReview_interactive_no_body(t *testing.T) {
 
 	shared.RunCommandFinder("", &api.PullRequest{ID: "THE-ID", Number: 123}, ghrepo.New("OWNER", "REPO"))
 
-	as, teardown := prompt.NewAskStubber()
-	defer teardown(t)
+	as := prompt.NewAskStubber(t)
 
 	as.StubPrompt("What kind of review do you want to give?").AnswerWith("Request changes")
 	as.StubPrompt("Review body").AnswerWith("")
diff --git a/pkg/cmd/workflow/enable/enable_test.go b/pkg/cmd/workflow/enable/enable_test.go
@@ -277,8 +277,7 @@ func TestEnableRun(t *testing.T) {
 		}
 
 		t.Run(tt.name, func(t *testing.T) {
-			as, teardown := prompt.NewAskStubber()
-			defer teardown(t)
+			as := prompt.NewAskStubber(t)
 			if tt.askStubs != nil {
 				tt.askStubs(as)
 			}
diff --git a/pkg/prompt/stubber.go b/pkg/prompt/stubber.go
@@ -15,9 +15,24 @@ type AskStubber struct {
 
 type testing interface {
 	Errorf(format string, args ...interface{})
+	Cleanup(func())
 }
 
-func NewAskStubber() (*AskStubber, func(t testing)) {
+func NewAskStubber(t testing) *AskStubber {
+	as, teardown := InitAskStubber()
+	t.Cleanup(func() {
+		teardown()
+		for _, s := range as.stubs {
+			if !s.matched {
+				t.Errorf("unmatched prompt stub: %+v", s)
+			}
+		}
+	})
+	return as
+}
+
+// Deprecated: use NewAskStubber
+func InitAskStubber() (*AskStubber, func()) {
 	origSurveyAsk := SurveyAsk
 	origSurveyAskOne := SurveyAskOne
 	as := AskStubber{}
@@ -124,29 +139,13 @@ func NewAskStubber() (*AskStubber, func(t testing)) {
 		return nil
 	}
 
-	teardown := func(t testing) {
+	teardown := func() {
 		SurveyAsk = origSurveyAsk
 		SurveyAskOne = origSurveyAskOne
-		for _, s := range as.stubs {
-			if !s.matched {
-				if t == nil {
-					panic(fmt.Sprintf("unmatched prompt stub: %+v", s))
-				} else {
-					t.Errorf("unmatched prompt stub: %+v", s)
-				}
-			}
-		}
 	}
 	return &as, teardown
 }
 
-func InitAskStubber() (*AskStubber, func()) {
-	as, teardown := NewAskStubber()
-	return as, func() {
-		teardown(nil)
-	}
-}
-
 type QuestionStub struct {
 	Name    string
 	Value   interface{}

Original file line number	Diff line number	Diff line change
`@@ -355,8 +355,7 @@ func Test_viewRun(t *testing.T) {`
`355`	`355`	`)),`
`356`	`356`	`)`
`357`	`357`
`358`		`- as, surveyteardown := prompt.NewAskStubber()`
`359`		`- defer surveyteardown(t)`
	`358`	`+ as := prompt.NewAskStubber(t)`
`360`	`359`	`as.StubPrompt("Select a gist").AnswerDefault()`
`361`	`360`	`}`
`362`	`361`
`@@ -470,8 +469,7 @@ func Test_promptGists(t *testing.T) {`
`470`	`469`	`client := &http.Client{Transport: reg}`
`471`	`470`
`472`	`471`	`t.Run(tt.name, func(t *testing.T) {`
`473`		`- as, surveyteardown := prompt.NewAskStubber()`
`474`		`- defer surveyteardown(t)`
	`472`	`+ as := prompt.NewAskStubber(t)`
`475`	`473`	`if tt.askStubs != nil {`
`476`	`474`	`tt.askStubs(as)`
`477`	`475`	`}`
Original file line number	Diff line number	Diff line change
`@@ -277,8 +277,7 @@ func TestEnableRun(t *testing.T) {`
`277`	`277`	`}`
`278`	`278`
`279`	`279`	`t.Run(tt.name, func(t *testing.T) {`
`280`		`- as, teardown := prompt.NewAskStubber()`
`281`		`- defer teardown(t)`
	`280`	`+ as := prompt.NewAskStubber(t)`
`282`	`281`	`if tt.askStubs != nil {`
`283`	`282`	`tt.askStubs(as)`
`284`	`283`	`}`