22package teatest
33
44import (
5- "bytes"
65 "fmt"
7- "io"
86 "os"
97 "os/signal"
108 "sync"
@@ -14,6 +12,7 @@ import (
1412
1513 tea "github.com/charmbracelet/bubbletea/v2"
1614 "github.com/charmbracelet/x/exp/golden"
15+ "github.com/charmbracelet/x/vt"
1716)
1817
1918// Program defines the subset of the tea.Program API we need for testing.
@@ -63,22 +62,22 @@ func WithDuration(d time.Duration) WaitForOption {
6362 }
6463}
6564
66- // WaitFor keeps reading from r until the condition matches.
65+ // WaitForOutput keeps reading from r until the condition matches.
6766// Default duration is 1s, default check interval is 50ms.
6867// These defaults can be changed with WithDuration and WithCheckInterval.
69- func WaitFor (
68+ func WaitForOutput (
7069 tb testing.TB ,
71- r io. Reader ,
72- condition func (bts [] byte ) bool ,
70+ tm * TestModel ,
71+ condition func (string ) bool ,
7372 options ... WaitForOption ,
7473) {
7574 tb .Helper ()
76- if err := doWaitFor (r , condition , options ... ); err != nil {
75+ if err := doWaitFor (tm , condition , options ... ); err != nil {
7776 tb .Fatal (err )
7877 }
7978}
8079
81- func doWaitFor (r io. Reader , condition func (bts [] byte ) bool , options ... WaitForOption ) error {
80+ func doWaitFor (tm * TestModel , condition func (string ) bool , options ... WaitForOption ) error {
8281 wf := WaitingForContext {
8382 Duration : time .Second ,
8483 CheckInterval : 50 * time .Millisecond , //nolint: gomnd
@@ -88,26 +87,21 @@ func doWaitFor(r io.Reader, condition func(bts []byte) bool, options ...WaitForO
8887 opt (& wf )
8988 }
9089
91- var b bytes.Buffer
9290 start := time .Now ()
9391 for time .Since (start ) <= wf .Duration {
94- if _ , err := io .ReadAll (io .TeeReader (r , & b )); err != nil {
95- return fmt .Errorf ("WaitFor: %w" , err )
96- }
97- if condition (b .Bytes ()) {
92+ if condition (tm .Output ()) {
9893 return nil
9994 }
10095 time .Sleep (wf .CheckInterval )
10196 }
102- return fmt .Errorf ("WaitFor: condition not met after %s. Last output:\n %s " , wf .Duration , b . String ())
97+ return fmt .Errorf ("WaitFor: condition not met after %s. Last output:\n %q " , wf .Duration , tm . Output ())
10398}
10499
105100// TestModel is a model that is being tested.
106101type TestModel struct {
107102 program * tea.Program
108103
109- in * bytes.Buffer
110- out io.ReadWriter
104+ term * vt.Terminal
111105
112106 modelCh chan tea.Model
113107 model tea.Model
@@ -118,17 +112,24 @@ type TestModel struct {
118112
119113// NewTestModel makes a new TestModel which can be used for tests.
120114func NewTestModel (tb testing.TB , m tea.Model , options ... TestOption ) * TestModel {
115+ var opts TestModelOptions
116+ for _ , opt := range options {
117+ opt (& opts )
118+ }
119+ if opts .size .Width == 0 {
120+ opts .size .Width , opts .size .Height = 70 , 40
121+ }
122+
121123 tm := & TestModel {
122- in : bytes .NewBuffer (nil ),
123- out : safe (bytes .NewBuffer (nil )),
124+ term : vt .NewTerminal (opts .size .Width , opts .size .Height ),
124125 modelCh : make (chan tea.Model , 1 ),
125126 doneCh : make (chan bool , 1 ),
126127 }
127128
128129 tm .program = tea .NewProgram (
129130 m ,
130- tea .WithInput (tm .in ),
131- tea .WithOutput (tm .out ),
131+ tea .WithInput (tm .term ),
132+ tea .WithOutput (tm .term ),
132133 tea .WithoutSignals (),
133134 )
134135
@@ -149,14 +150,7 @@ func NewTestModel(tb testing.TB, m tea.Model, options ...TestOption) *TestModel
149150 tm .program .Kill ()
150151 }()
151152
152- var opts TestModelOptions
153- for _ , opt := range options {
154- opt (& opts )
155- }
156-
157- if opts .size .Width != 0 {
158- tm .program .Send (opts .size )
159- }
153+ tm .program .Send (opts .size )
160154 return tm
161155}
162156
@@ -229,14 +223,14 @@ func (tm *TestModel) FinalModel(tb testing.TB, opts ...FinalOpt) tea.Model {
229223// FinalOutput returns the program's final output io.Reader.
230224// This method only returns once the program has finished running or when it
231225// times out.
232- func (tm * TestModel ) FinalOutput (tb testing.TB , opts ... FinalOpt ) io. Reader {
226+ func (tm * TestModel ) FinalOutput (tb testing.TB , opts ... FinalOpt ) string {
233227 tm .waitDone (tb , opts )
234228 return tm .Output ()
235229}
236230
237231// Output returns the program's current output io.Reader.
238- func (tm * TestModel ) Output () io. Reader {
239- return tm .out
232+ func (tm * TestModel ) Output () string {
233+ return tm .term . String ()
240234}
241235
242236// Send sends messages to the underlying program.
@@ -271,31 +265,7 @@ func (tm *TestModel) GetProgram() *tea.Program {
271265// Important: this uses the system `diff` tool.
272266//
273267// You can update the golden files by running your tests with the -update flag.
274- func RequireEqualOutput (tb testing.TB , out [] byte ) {
268+ func RequireEqualOutput (tb testing.TB , out string ) {
275269 tb .Helper ()
276- golden .RequireEqualEscape (tb , out , true )
277- }
278-
279- func safe (rw io.ReadWriter ) io.ReadWriter {
280- return & safeReadWriter {rw : rw }
281- }
282-
283- // safeReadWriter implements io.ReadWriter, but locks reads and writes.
284- type safeReadWriter struct {
285- rw io.ReadWriter
286- m sync.RWMutex
287- }
288-
289- // Read implements io.ReadWriter.
290- func (s * safeReadWriter ) Read (p []byte ) (n int , err error ) {
291- s .m .RLock ()
292- defer s .m .RUnlock ()
293- return s .rw .Read (p ) //nolint: wrapcheck
294- }
295-
296- // Write implements io.ReadWriter.
297- func (s * safeReadWriter ) Write (p []byte ) (int , error ) {
298- s .m .Lock ()
299- defer s .m .Unlock ()
300- return s .rw .Write (p ) //nolint: wrapcheck
270+ golden .RequireEqualEscape (tb , []byte (out ), true )
301271}
0 commit comments