feat: FinalOuput and FinalModel allow a timeout (#2)

* feat: FinalOuput and FinalModel allow a timeout Signed-off-by: Carlos Alexandro Becker <[email protected]> * docs: godoc Signed-off-by: Carlos Alexandro Becker <[email protected]> * feat: add a public WaitFinished method too Signed-off-by: Carlos Alexandro Becker <[email protected]> --------- Signed-off-by: Carlos Alexandro Becker <[email protected]>
charmbracelet · May 8, 2023 · 2bd6fa1 · 2bd6fa1
1 parent 0a20c60
commit 2bd6fa1
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 11 deletions.
diff --git a/exp/teatest/app_test.go b/exp/teatest/app_test.go
@@ -35,13 +35,13 @@ func TestApp(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	out := readBts(t, tm.FinalOutput())
+	out := readBts(t, tm.FinalOutput(t, teatest.WithFinalTimeout(time.Second)))
 	if !regexp.MustCompile(`This program will exit in \d+ seconds`).Match(out) {
 		t.Fatalf("output does not match the given regular expression: %s", string(out))
 	}
 	teatest.RequireEqualOutput(t, out)
 
-	if tm.FinalModel().(model) != 9 {
+	if tm.FinalModel(t).(model) != 9 {
 		t.Errorf("expected model to be 10, was %d", m)
 	}
 }
@@ -72,7 +72,7 @@ func TestAppInteractive(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	if tm.FinalModel().(model) != 7 {
+	if tm.FinalModel(t).(model) != 7 {
 		t.Errorf("expected model to be 7, was %d", m)
 	}
 }

diff --git a/exp/teatest/teatest.go b/exp/teatest/teatest.go
@@ -162,16 +162,52 @@ func NewTestModel(tb testing.TB, m tea.Model, options ...TestOption) *TestModel
 	return tm
 }
 
-func (tm *TestModel) waitDone() {
+func (tm *TestModel) waitDone(tb testing.TB, opts []FinalOpt) {
 	tm.done.Do(func() {
-		<-tm.doneCh
+		fopts := FinalOpts{}
+		for _, opt := range opts {
+			opt(&fopts)
+		}
+		if fopts.timeout > 0 {
+			select {
+			case <-time.After(fopts.timeout):
+				tb.Fatalf("timeout after %s", fopts.timeout)
+			case <-tm.doneCh:
+			}
+		} else {
+			<-tm.doneCh
+		}
 	})
 }
 
+// FinalOpts represents the options for FinalModel and FinalOutput.
+type FinalOpts struct {
+	timeout time.Duration
+}
+
+// FinalOpt changes FinalOpts.
+type FinalOpt func(opts *FinalOpts)
+
+// WithFinalTimeout allows to set a timeout for how long FinalModel and
+// FinalOuput should wait for the program to complete.
+func WithFinalTimeout(d time.Duration) FinalOpt {
+	return func(opts *FinalOpts) {
+		opts.timeout = d
+	}
+}
+
+// WaitFinished waits for the app to finish.
+// This method only returns once the program has finished running or when it
+// times out.
+func (tm *TestModel) WaitFinished(tb testing.TB, opts ...FinalOpt) {
+	tm.waitDone(tb, opts)
+}
+
 // FinalModel returns the resulting model, resulting from program.Run().
-// This method only returns once the program has finished running.
-func (tm *TestModel) FinalModel() tea.Model {
-	tm.waitDone()
+// This method only returns once the program has finished running or when it
+// times out.
+func (tm *TestModel) FinalModel(tb testing.TB, opts ...FinalOpt) tea.Model {
+	tm.waitDone(tb, opts)
 	select {
 	case m := <-tm.modelCh:
 		tm.model = m
@@ -182,9 +218,10 @@ func (tm *TestModel) FinalModel() tea.Model {
 }
 
 // FinalOutput returns the program's final output io.Reader.
-// It'll block until the program finishes.
-func (tm *TestModel) FinalOutput() io.Reader {
-	tm.waitDone()
+// This method only returns once the program has finished running or when it
+// times out.
+func (tm *TestModel) FinalOutput(tb testing.TB, opts ...FinalOpt) io.Reader {
+	tm.waitDone(tb, opts)
 	return tm.Output()
 }