Skip to content

Commit 34a5480

Browse files
committed
update RunWithStatus(...) to respect custom writer attached by WithContextWriter(...)
Signed-off-by: Tiger Wang <tigerwang@outlook.com>
1 parent e42673d commit 34a5480

File tree

2 files changed

+72
-6
lines changed

2 files changed

+72
-6
lines changed

pkg/progress/writer.go

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -66,17 +66,27 @@ func Run(ctx context.Context, pf progressFunc) error {
6666
// RunWithStatus will run a writer and the progress function in parallel and return a status
6767
func RunWithStatus(ctx context.Context, pf progressFuncWithStatus) (string, error) {
6868
eg, _ := errgroup.WithContext(ctx)
69-
w, err := NewWriter(os.Stderr)
70-
var result string
71-
if err != nil {
72-
return "", err
69+
70+
var (
71+
w Writer
72+
err error
73+
result string
74+
)
75+
76+
if ctx.Value(writerKey{}) != nil {
77+
w = ctx.Value(writerKey{}).(Writer)
78+
} else {
79+
w, err = NewWriter(os.Stderr)
80+
if err != nil {
81+
return "", err
82+
}
83+
ctx = WithContextWriter(ctx, w)
7384
}
85+
7486
eg.Go(func() error {
7587
return w.Start(context.Background())
7688
})
7789

78-
ctx = WithContextWriter(ctx, w)
79-
8090
eg.Go(func() error {
8191
defer w.Stop()
8292
s, err := pf(ctx)

pkg/progress/writer_test.go

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ package progress
1818

1919
import (
2020
"context"
21+
"os"
22+
"strings"
2123
"testing"
2224

2325
"gotest.tools/v3/assert"
@@ -29,3 +31,57 @@ func TestNoopWriter(t *testing.T) {
2931

3032
assert.Equal(t, writer, &noopWriter{})
3133
}
34+
35+
func TestRunWithStatusWithoutCustomContextWriter(t *testing.T) {
36+
r, w, err := os.Pipe()
37+
assert.NilError(t, err)
38+
39+
os.Stderr = w // mock Stderr for default writer just for testing purpose
40+
41+
result := make(chan string)
42+
go func() {
43+
buf := make([]byte, 256)
44+
n, _ := r.Read(buf)
45+
result <- string(buf[:n])
46+
}()
47+
48+
// run without any custom writer, so it will use the default writer
49+
_, err = RunWithStatus(context.TODO(), func(ctx context.Context) (string, error) {
50+
ContextWriter(ctx).Event(Event{Text: "pass"})
51+
return "test", nil
52+
})
53+
54+
assert.NilError(t, err)
55+
56+
actual := <-result
57+
assert.Equal(t, strings.TrimSpace(actual), "pass")
58+
}
59+
60+
func TestRunWithStatusrWithCustomContextWriter(t *testing.T) {
61+
r, w, err := os.Pipe()
62+
assert.NilError(t, err)
63+
64+
writer, err := NewWriter(w) // custom writer
65+
assert.NilError(t, err)
66+
67+
result := make(chan string)
68+
go func() {
69+
buf := make([]byte, 256)
70+
n, _ := r.Read(buf)
71+
result <- string(buf[:n])
72+
}()
73+
74+
// attach the custom writer to the context
75+
ctx := WithContextWriter(context.TODO(), writer)
76+
77+
// run with the custom writer
78+
_, err = RunWithStatus(ctx, func(ctx context.Context) (string, error) {
79+
ContextWriter(ctx).Event(Event{Text: "pass"})
80+
return "test", nil
81+
})
82+
83+
assert.NilError(t, err)
84+
85+
actual := <-result
86+
assert.Equal(t, strings.TrimSpace(actual), "pass")
87+
}

0 commit comments

Comments
 (0)