diff --git a/docs/pkg/cli/streaming.md b/docs/pkg/cli/streaming.md index 67dd41b..25d0d31 100644 --- a/docs/pkg/cli/streaming.md +++ b/docs/pkg/cli/streaming.md @@ -34,7 +34,8 @@ When word-wrap is enabled, the stream tracks the current column position and ins ## Custom Output Writer -By default, streams write to `os.Stdout`. Redirect to any `io.Writer`: +By default, streams write to the CLI stdout writer (`stdoutWriter()`), so tests can +redirect output via `cli.SetStdout` and other callers can provide any `io.Writer`: ```go var buf strings.Builder @@ -77,7 +78,7 @@ stream.Done() | Option | Description | |--------|-------------| | `WithWordWrap(cols)` | Set the word-wrap column width | -| `WithStreamOutput(w)` | Set the output writer (default: `os.Stdout`) | +| `WithStreamOutput(w)` | Set the output writer (default: `stdoutWriter()`) | ## Example: LLM Token Streaming diff --git a/pkg/cli/stream.go b/pkg/cli/stream.go index 66de4fd..a324bba 100644 --- a/pkg/cli/stream.go +++ b/pkg/cli/stream.go @@ -3,7 +3,6 @@ package cli import ( "fmt" "io" - "os" "strings" "sync" @@ -12,7 +11,8 @@ import ( // StreamOption configures a Stream. // -// stream := cli.NewStream(cli.WithWordWrap(80), cli.WithStreamOutput(os.Stdout)) +// stream := cli.NewStream(cli.WithWordWrap(80)) +// stream.Wait() type StreamOption func(*Stream) // WithWordWrap sets the word-wrap column width. @@ -20,7 +20,7 @@ func WithWordWrap(cols int) StreamOption { return func(s *Stream) { s.wrap = cols } } -// WithStreamOutput sets the output writer (default: os.Stdout). +// WithStreamOutput sets the output writer (default: stdoutWriter()). func WithStreamOutput(w io.Writer) StreamOption { return func(s *Stream) { s.out = w } } @@ -48,7 +48,7 @@ type Stream struct { // NewStream creates a streaming text renderer. func NewStream(opts ...StreamOption) *Stream { s := &Stream{ - out: os.Stdout, + out: stdoutWriter(), done: make(chan struct{}), } for _, opt := range opts { diff --git a/pkg/cli/stream_test.go b/pkg/cli/stream_test.go index 589162f..86ff5fd 100644 --- a/pkg/cli/stream_test.go +++ b/pkg/cli/stream_test.go @@ -10,6 +10,19 @@ import ( ) func TestStream_Good(t *testing.T) { + t.Run("uses injected stdout by default", func(t *testing.T) { + var buf bytes.Buffer + SetStdout(&buf) + defer SetStdout(nil) + s := NewStream() + + s.Write("hello") + s.Done() + s.Wait() + + assert.Equal(t, "hello\n", buf.String()) + }) + t.Run("basic write", func(t *testing.T) { var buf bytes.Buffer s := NewStream(WithStreamOutput(&buf))