diff --git a/pkg/cli/stream.go b/pkg/cli/stream.go index 42c5aed..d78a442 100644 --- a/pkg/cli/stream.go +++ b/pkg/cli/stream.go @@ -39,6 +39,7 @@ type Stream struct { wrap int col int // current column position (visible characters) done chan struct{} + once sync.Once mu sync.Mutex } @@ -107,12 +108,14 @@ func (s *Stream) WriteFrom(r io.Reader) error { // Done signals that no more text will arrive. func (s *Stream) Done() { - s.mu.Lock() - if s.col > 0 { - fmt.Fprintln(s.out) // ensure trailing newline - } - s.mu.Unlock() - close(s.done) + s.once.Do(func() { + s.mu.Lock() + if s.col > 0 { + fmt.Fprintln(s.out) // ensure trailing newline + } + s.mu.Unlock() + close(s.done) + }) } // Wait blocks until Done is called. diff --git a/pkg/cli/stream_test.go b/pkg/cli/stream_test.go index 6937547..3a528ce 100644 --- a/pkg/cli/stream_test.go +++ b/pkg/cli/stream_test.go @@ -153,6 +153,18 @@ func TestStream_Good(t *testing.T) { assert.Equal(t, "text\n", buf.String()) // no double newline }) + t.Run("Done is idempotent", func(t *testing.T) { + var buf bytes.Buffer + s := NewStream(WithStreamOutput(&buf)) + + s.Write("text") + s.Done() + s.Done() + s.Wait() + + assert.Equal(t, "text\n", buf.String()) + }) + t.Run("word wrap uses visible width", func(t *testing.T) { var buf bytes.Buffer s := NewStream(WithWordWrap(4), WithStreamOutput(&buf))