feat: add WaitForMsg function

exp/teatest/app_test.go

@@ -60,7 +60,7 @@ func TestAppInteractive(t *testing.T) {
 		t.Fatalf("output does not match: expected %q", string(bts))
 	}
 
-	teatest.WaitFor(t, tm.Output(), func(out []byte) bool {
+	teatest.WaitForOutput(t, tm.Output(), func(out []byte) bool {
 		return bytes.Contains(out, []byte("This program will exit in 7 seconds"))
 	}, teatest.WithDuration(5*time.Second), teatest.WithCheckInterval(time.Millisecond*10))
 

exp/teatest/msg_buffer.go

@@ -0,0 +1,31 @@
+package teatest
+
+import (
+	"sync"
+
+	tea "github.com/charmbracelet/bubbletea"
+)
+
+// msgBuffer stores messages for checking in WaitForMsg.
+type msgBuffer struct {
+	msgs []tea.Msg
+	mu   sync.Mutex
+}
+
+func (b *msgBuffer) append(msg tea.Msg) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	b.msgs = append(b.msgs, msg)
+}
+
+// forEach executes the given function for each message while holding the lock.
+func (b *msgBuffer) forEach(fn func(msg tea.Msg) bool) tea.Msg {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	for _, msg := range b.msgs {
+		if fn(msg) {
+			return msg
+		}
+	}
+	return nil
+}

exp/teatest/msg_capture.go

@@ -0,0 +1,32 @@
+package teatest
+
+import (
+	tea "github.com/charmbracelet/bubbletea"
+)
+
+// msgCaptureModel wraps a model to capture messages.
+type msgCaptureModel struct {
+	model  tea.Model
+	buffer *msgBuffer
+}
+
+func (m msgCaptureModel) Init() tea.Cmd {
+	return m.model.Init()
+}
+
+func (m msgCaptureModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	m.buffer.append(msg)
+	model, cmd := m.model.Update(msg)
+	if wrappedModel, ok := model.(msgCaptureModel); ok {
+		return wrappedModel, cmd
+	}
+
+	return msgCaptureModel{
+		model:  model,
+		buffer: m.buffer,
+	}, cmd
+}
+
+func (m msgCaptureModel) View() string {
+	return m.model.View()
+}

exp/teatest/teatest.go

@@ -63,22 +63,22 @@ func WithDuration(d time.Duration) WaitForOption {
 	}
 }
 
-// WaitFor keeps reading from r until the condition matches.
+// WaitForOutput keeps reading from r until the condition matches.
 // Default duration is 1s, default check interval is 50ms.
 // These defaults can be changed with WithDuration and WithCheckInterval.
-func WaitFor(
+func WaitForOutput(
 	tb testing.TB,
 	r io.Reader,
 	condition func(bts []byte) bool,
 	options ...WaitForOption,
 ) {
 	tb.Helper()
-	if err := doWaitFor(r, condition, options...); err != nil {
+	if err := doWaitForOutput(r, condition, options...); err != nil {
 		tb.Fatal(err)
 	}
 }
 
-func doWaitFor(r io.Reader, condition func(bts []byte) bool, options ...WaitForOption) error {
+func doWaitForOutput(r io.Reader, condition func(bts []byte) bool, options ...WaitForOption) error {
 	wf := WaitingForContext{
 		Duration:      time.Second,
 		CheckInterval: 50 * time.Millisecond, //nolint: mnd
@@ -114,6 +114,8 @@ type TestModel struct {
 
 	done   sync.Once
 	doneCh chan bool
+
+	msgs *msgBuffer
 }
 
 // NewTestModel makes a new TestModel which can be used for tests.
@@ -123,11 +125,19 @@ func NewTestModel(tb testing.TB, m tea.Model, options ...TestOption) *TestModel
 		out:     safe(bytes.NewBuffer(nil)),
 		modelCh: make(chan tea.Model, 1),
 		doneCh:  make(chan bool, 1),
+		msgs: &msgBuffer{
+			msgs: make([]tea.Msg, 0),
+		},
+	}
+
+	wrappedModel := msgCaptureModel{
+		model:  m,
+		buffer: tm.msgs,
 	}
 
 	//nolint: staticcheck
 	tm.program = tea.NewProgram(
-		m,
+		wrappedModel,
 		tea.WithInput(tm.in),
 		tea.WithOutput(tm.out),
 		tea.WithoutSignals(),
@@ -162,6 +172,45 @@ func NewTestModel(tb testing.TB, m tea.Model, options ...TestOption) *TestModel
 	return tm
 }
 
+// WaitForMsg keeps checking messages until the condition matches or timeout is reached.
+// Default duration is 1s, default check interval is 50ms.
+func (tm *TestModel) WaitForMsg(
+	tb testing.TB,
+	condition func(msg tea.Msg) bool,
+	options ...WaitForOption,
+) tea.Msg {
+	tb.Helper()
+	msg, err := tm.doWaitForMsg(condition, options...)
+	if err != nil {
+		tb.Fatal(err)
+	}
+	return msg
+}
+
+func (tm *TestModel) doWaitForMsg(
+	condition func(msg tea.Msg) bool,
+	options ...WaitForOption,
+) (tea.Msg, error) {
+	wf := WaitingForContext{
+		Duration:      time.Second,
+		CheckInterval: 50 * time.Millisecond,
+	}
+
+	for _, opt := range options {
+		opt(&wf)
+	}
+
+	start := time.Now()
+	for time.Since(start) <= wf.Duration {
+		if msg := tm.msgs.forEach(condition); msg != nil {
+			return msg, nil
+		}
+		time.Sleep(wf.CheckInterval)
+	}
+
+	return nil, fmt.Errorf("WaitForMsg: condition not met after %s", wf.Duration)
+}
+
 func (tm *TestModel) waitDone(tb testing.TB, opts []FinalOpt) {
 	tm.done.Do(func() {
 		fopts := FinalOpts{}

exp/teatest/teatest_test.go

@@ -11,7 +11,7 @@ import (
 )
 
 func TestWaitForErrorReader(t *testing.T) {
-	err := doWaitFor(iotest.ErrReader(fmt.Errorf("fake")), func(bts []byte) bool {
+	err := doWaitForOutput(iotest.ErrReader(fmt.Errorf("fake")), func(bts []byte) bool {
 		return true
 	}, WithDuration(time.Millisecond), WithCheckInterval(10*time.Microsecond))
 	if err == nil {
@@ -23,7 +23,7 @@ func TestWaitForErrorReader(t *testing.T) {
 }
 
 func TestWaitForTimeout(t *testing.T) {
-	err := doWaitFor(strings.NewReader("nope"), func(bts []byte) bool {
+	err := doWaitForOutput(strings.NewReader("nope"), func(bts []byte) bool {
 		return false
 	}, WithDuration(time.Millisecond), WithCheckInterval(10*time.Microsecond))
 	if err == nil {