update RunWithStatus(...) to respect custom writer attached by WithContextWriter(...)

pkg/progress/writer.go

@@ -76,17 +76,27 @@ func RunWithTitle(ctx context.Context, pf progressFunc, out io.Writer, progressT
 // RunWithStatus will run a writer and the progress function in parallel and return a status
 func RunWithStatus(ctx context.Context, pf progressFuncWithStatus, out io.Writer, progressTitle string) (string, error) {
 	eg, _ := errgroup.WithContext(ctx)
+
 	w, err := NewWriter(ctx, out, progressTitle)
 	var result string
 	if err != nil {
 		return "", err
+
 	}
+
+	if w == nil {
+		_w, err := NewWriter(os.Stderr)
+		if err != nil {
+			return "", err
+		}
+		w = _w
+		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)
@@ -96,7 +106,7 @@ func RunWithStatus(ctx context.Context, pf progressFuncWithStatus, out io.Writer
 		return err
 	})
 
-	err = eg.Wait()
+	err := eg.Wait()
 	return result, err
 }
 

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")
+}