From f240be0daccaf6fac5473599e5c296e73437a415 Mon Sep 17 00:00:00 2001 From: Broderick Westrope Date: Wed, 25 Dec 2024 20:06:50 +1100 Subject: [PATCH 1/3] refactor: rename WaitForOutput function --- exp/teatest/app_test.go | 2 +- exp/teatest/teatest.go | 8 ++++---- exp/teatest/teatest_test.go | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/exp/teatest/app_test.go b/exp/teatest/app_test.go index 36712356..7ab06e46 100644 --- a/exp/teatest/app_test.go +++ b/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)) diff --git a/exp/teatest/teatest.go b/exp/teatest/teatest.go index 8308e82b..8ed9f66b 100644 --- a/exp/teatest/teatest.go +++ b/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 diff --git a/exp/teatest/teatest_test.go b/exp/teatest/teatest_test.go index c83f484f..a9cd4def 100644 --- a/exp/teatest/teatest_test.go +++ b/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 { From ade9ebad530e832d0d9d956c0770e872a8984d79 Mon Sep 17 00:00:00 2001 From: Broderick Westrope Date: Wed, 25 Dec 2024 20:15:27 +1100 Subject: [PATCH 2/3] feat: wrap inner model to capture all messages --- exp/teatest/msg_buffer.go | 31 +++++++++++++++++++++++++++++++ exp/teatest/msg_capture.go | 32 ++++++++++++++++++++++++++++++++ exp/teatest/teatest.go | 12 +++++++++++- 3 files changed, 74 insertions(+), 1 deletion(-) create mode 100644 exp/teatest/msg_buffer.go create mode 100644 exp/teatest/msg_capture.go diff --git a/exp/teatest/msg_buffer.go b/exp/teatest/msg_buffer.go new file mode 100644 index 00000000..b8c36ec1 --- /dev/null +++ b/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 +} diff --git a/exp/teatest/msg_capture.go b/exp/teatest/msg_capture.go new file mode 100644 index 00000000..b8fba1eb --- /dev/null +++ b/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() +} diff --git a/exp/teatest/teatest.go b/exp/teatest/teatest.go index 8ed9f66b..c733bd14 100644 --- a/exp/teatest/teatest.go +++ b/exp/teatest/teatest.go @@ -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(), From e76c1e68d13c9d8246e2c14f6f3399e323cadc51 Mon Sep 17 00:00:00 2001 From: Broderick Westrope Date: Wed, 25 Dec 2024 20:15:47 +1100 Subject: [PATCH 3/3] feat: add WaitForMsg function --- exp/teatest/teatest.go | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/exp/teatest/teatest.go b/exp/teatest/teatest.go index c733bd14..bcce87be 100644 --- a/exp/teatest/teatest.go +++ b/exp/teatest/teatest.go @@ -172,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{}