44 "bytes"
55 "io"
66 "os"
7+ "sync"
78)
89
910// Sequence represents an ANSI escape sequence.
@@ -61,7 +62,7 @@ func (h *Handle) Stdout() *Printer {
6162// Printer allows for writing data with optional ANSI escape sequences based on
6263// the color settings for a target.
6364type Printer struct {
64- file * os. File
65+ file io. Writer
6566 buf bytes.Buffer
6667 useColor bool
6768}
@@ -81,6 +82,25 @@ func newPrinter(file *os.File, isTerm bool, c Color) *Printer {
8182 return & Printer {file : file , useColor : useColor }
8283}
8384
85+ // TestPrinter returns a Printer suitable for testing. All output, including
86+ // flushed data, is captured and accessible via Bytes.
87+ func TestPrinter (useColor bool ) * Printer {
88+ return & Printer {file : & lockedBuffer {}, useColor : useColor }
89+ }
90+
91+ // lockedBuffer is a goroutine-safe bytes.Buffer used as the flush target in
92+ // test printers so that Bytes can return the combined buffered + flushed data.
93+ type lockedBuffer struct {
94+ mu sync.Mutex
95+ buf bytes.Buffer
96+ }
97+
98+ func (lb * lockedBuffer ) Write (p []byte ) (int , error ) {
99+ lb .mu .Lock ()
100+ defer lb .mu .Unlock ()
101+ return lb .buf .Write (p )
102+ }
103+
84104// Set writes the provided Sequence.
85105func (p * Printer ) Set (s Sequence ) {
86106 if p .useColor {
@@ -108,8 +128,17 @@ func (p *Printer) Discard() {
108128 p .buf .Reset ()
109129}
110130
111- // Bytes returns the current contents of the buffer.
131+ // Bytes returns the current contents of the buffer. For test printers created
132+ // with TestPrinter, this also includes previously flushed data.
112133func (p * Printer ) Bytes () []byte {
134+ if lb , ok := p .file .(* lockedBuffer ); ok {
135+ lb .mu .Lock ()
136+ flushed := lb .buf .Bytes ()
137+ lb .mu .Unlock ()
138+ if len (flushed ) > 0 {
139+ return append (flushed , p .buf .Bytes ()... )
140+ }
141+ }
113142 return p .buf .Bytes ()
114143}
115144
0 commit comments