update RunWithStatus(...) to respect custom writer attached by WithCo…

…ntextWriter(...)

Signed-off-by: Tiger Wang <tigerwang@outlook.com>

pkg/progress/writer.go

@@ -66,17 +66,27 @@ func Run(ctx context.Context, pf progressFunc) error {
 // RunWithStatus will run a writer and the progress function in parallel and return a status
 func RunWithStatus(ctx context.Context, pf progressFuncWithStatus) (string, error) {
 	eg, _ := errgroup.WithContext(ctx)
-	w, err := NewWriter(os.Stderr)
-	var result string
-	if err != nil {
-		return "", err
+
+	var (
+		w      Writer
+		err    error
+		result string
+	)
+
+	if ctx.Value(writerKey{}) != nil {
+		w = ctx.Value(writerKey{}).(Writer)
+	} else {
+		w, err = NewWriter(os.Stderr)
+		if err != nil {
+			return "", err
+		}
+		ctx = WithContextWriter(ctx, w)
 	}
+
 	eg.Go(func() error {
 		return w.Start(context.Background())
 	})
 
-	ctx = WithContextWriter(ctx, w)
-
 	eg.Go(func() error {
 		defer w.Stop()
 		s, err := pf(ctx)

pkg/progress/writer_test.go

@@ -18,6 +18,8 @@ package progress
 
 import (
 	"context"
+	"os"
+	"strings"
 	"testing"
 
 	"gotest.tools/v3/assert"
@@ -29,3 +31,57 @@ func TestNoopWriter(t *testing.T) {
 
 	assert.Equal(t, writer, &noopWriter{})
 }
+
+func TestRunWithStatusWithoutCustomContextWriter(t *testing.T) {
+	r, w, err := os.Pipe()
+	assert.NilError(t, err)
+
+	os.Stderr = w // mock Stderr for default writer just for testing purpose
+
+	result := make(chan string)
+	go func() {
+		buf := make([]byte, 256)
+		n, _ := r.Read(buf)
+		result <- string(buf[:n])
+	}()
+
+	// run without any custom writer, so it will use the default writer
+	_, err = RunWithStatus(context.TODO(), func(ctx context.Context) (string, error) {
+		ContextWriter(ctx).Event(Event{Text: "pass"})
+		return "test", nil
+	})
+
+	assert.NilError(t, err)
+
+	actual := <-result
+	assert.Equal(t, strings.TrimSpace(actual), "pass")
+}
+
+func TestRunWithStatusrWithCustomContextWriter(t *testing.T) {
+	r, w, err := os.Pipe()
+	assert.NilError(t, err)
+
+	writer, err := NewWriter(w) // custom writer
+	assert.NilError(t, err)
+
+	result := make(chan string)
+	go func() {
+		buf := make([]byte, 256)
+		n, _ := r.Read(buf)
+		result <- string(buf[:n])
+	}()
+
+	// attach the custom writer to the context
+	ctx := WithContextWriter(context.TODO(), writer)
+
+	// run with the custom writer
+	_, err = RunWithStatus(ctx, func(ctx context.Context) (string, error) {
+		ContextWriter(ctx).Event(Event{Text: "pass"})
+		return "test", nil
+	})
+
+	assert.NilError(t, err)
+
+	actual := <-result
+	assert.Equal(t, strings.TrimSpace(actual), "pass")
+}