image-builder: reuse progress bars for upload

cmd/image-builder/main.go

@@ -689,6 +689,7 @@ func cmdBuild(cmd *cobra.Command, args []string) error {
 
 	pbar, err := progressFromCmd(cmd, progress.ProgressConfig{
 		FilePath: filepath.Join(outputDir, fmt.Sprintf("%s.progress", basenameFor(img, outputBasename))),
+		WithMsg:  true,
 	})
 	if err != nil {
 		return err
@@ -749,20 +750,30 @@ func cmdBuild(cmd *cobra.Command, args []string) error {
 		return err
 	}
 	pbar.Stop()
-
 	fmt.Fprintf(osStdout, "Image build successful: %s\n", imagePath)
 
+	pbar, err = progressFromCmd(cmd, progress.ProgressConfig{
+		FilePath: filepath.Join(outputDir, fmt.Sprintf("%s.progress", basenameFor(img, outputBasename))),
+		Bytes:    true,
+		Speed:    true,
+	})
+	if err != nil {
+		return err
+	}
 	// Default upload result to write out in case no uploader was specified
 	uploadResult := &cloud.UploadResult{
 		Provider: "LocalPath",
 		ImageID:  imagePath,
 	}
 	if uploader != nil {
+		pbar.Start()
+		pbar.SetPulseMsgf("Uploading")
 		// XXX: integrate better into the progress, see bib
-		uploadResult, err = uploadImageWithProgress(uploader, imagePath)
+		uploadResult, err = uploadImageWithProgress(uploader, pbar, imagePath)
 		if err != nil {
 			return err
 		}
+		pbar.Stop()
 	}
 	if withUploadResult {
 		p := filepath.Join(outputDir, fmt.Sprintf("%s.upload-result", basenameFor(img, outputBasename)))

cmd/image-builder/upload.go

@@ -10,7 +10,6 @@ import (
 	"path/filepath"
 	"strings"
 
-	"github.com/cheggaaa/pb/v3"
 	"github.com/spf13/cobra"
 	"go.yaml.in/yaml/v3"
 
@@ -44,7 +43,7 @@ var (
 	ibmNewUploader       = ibmcloud.NewUploader
 )
 
-func uploadImageWithProgress(uploader cloud.Uploader, imagePath string) (*cloud.UploadResult, error) {
+func uploadImageWithProgress(uploader cloud.Uploader, pbar progress.ProgressBar, imagePath string) (*cloud.UploadResult, error) {
 	f, err := os.Open(imagePath)
 	if err != nil {
 		return nil, err
@@ -60,13 +59,14 @@ func uploadImageWithProgress(uploader cloud.Uploader, imagePath string) (*cloud.
 	if sizei64 < 0 {
 		return nil, fmt.Errorf("invalid size read for %s: %d", imagePath, sizei64)
 	}
+	sizei := int(sizei64)
 	size := uint64(sizei64)
-	pbar := pb.New64(st.Size())
-	pbar.Set(pb.Bytes, true)
-	pbar.SetWriter(osStderr)
-	r := pbar.NewProxyReader(f)
+	r, err := progress.NewProxyReader(f, sizei, pbar)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create proxy reader: %w", err)
+	}
 	pbar.Start()
-	defer pbar.Finish()
+	pbar.SetPulseMsgf("Uploading step")
 
 	return uploader.UploadAndRegister(r, size, osStderr)
 }
@@ -372,7 +372,16 @@ func cmdUpload(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	result, err := uploadImageWithProgress(uploader, imagePath)
+	pbar, err := progress.New("auto", progress.ProgressConfig{
+		Bytes: true,
+		Speed: true,
+	})
+	if err != nil {
+		return err
+	}
+	pbar.Start()
+	defer pbar.Stop()
+	result, err := uploadImageWithProgress(uploader, pbar, imagePath)
 	if err != nil {
 		return err
 	}

cmd/image-builder/upload_test.go

@@ -3,6 +3,7 @@ package main_test
 import (
 	"bytes"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 	"strings"
@@ -89,9 +90,25 @@ func TestUploadWithAWSMock(t *testing.T) {
 		})
 		defer restore()
 
+		// to capture progress, overwrite os.Stderr entirely
+		old := os.Stderr
+		defer func() {
+			os.Stderr = old
+		}()
+		r, w, err := os.Pipe()
+		assert.NoError(t, err)
+		defer w.Close()
+		defer r.Close()
+		os.Stderr = w
+
 		err = main.Run()
 		require.NoError(t, err)
 
+		assert.NoError(t, w.Close())
+		var progressBuf bytes.Buffer
+		_, err = io.Copy(&progressBuf, r)
+		assert.NoError(t, err)
+
 		assert.Equal(t, regionName, "aws-region-1")
 		assert.Equal(t, bucketName, "aws-bucket-2")
 		assert.Equal(t, amiName, "aws-ami-3")
@@ -103,12 +120,11 @@ func TestUploadWithAWSMock(t *testing.T) {
 		assert.Equal(t, 0, fa.checkCalls)
 		assert.Equal(t, 1, fa.uploadAndRegisterCalls)
 		assert.Equal(t, fakeDiskContent, fa.uploadAndRegisterRead.String())
-		// progress was rendered to stderr
-		assert.Contains(t, fakeStderr.String(), "--] 100.00%")
-
 		// warning was passed to stderr
 		assert.Contains(t, fakeStderr.String(), tc.expectedWarning)
 
+		// progress was rendered to stderr
+		assert.Contains(t, progressBuf.String(), "Uploading step")
 	}
 }
 

pkg/progress/progress.go

@@ -83,6 +83,11 @@ type ProgressBar interface {
 type ProgressConfig struct {
 	// file progress only
 	FilePath string
+
+	// term progress only
+	Bytes   bool
+	Speed   bool
+	WithMsg bool
 }
 
 // New creates a new progressbar based on the requested type
@@ -95,17 +100,17 @@ func New(typ string, config ProgressConfig) (ProgressBar, error) {
 		// autoselect based on if we are on an interactive
 		// terminal, use verbose progress for scripts
 		if isattyIsTerminal(os.Stdin.Fd()) && w > 0 && h > 0 {
-			return NewTerminalProgressBar()
+			return NewTerminalProgressBar(&config)
 		}
 		return NewVerboseProgressBar()
 	case "verbose":
 		return NewVerboseProgressBar()
 	case "term":
-		return NewTerminalProgressBar()
+		return NewTerminalProgressBar(&config)
 	case "debug":
 		return NewDebugProgressBar()
 	case "file":
-		return NewFileProgressBar(config.FilePath)
+		return NewFileProgressBar(&config)
 	default:
 		return nil, fmt.Errorf("unknown progress type: %q", typ)
 	}
@@ -114,6 +119,9 @@ func New(typ string, config ProgressConfig) (ProgressBar, error) {
 type terminalProgressBar struct {
 	mu sync.Mutex
 
+	bytes bool
+	speed bool
+
 	spinnerPb   *pb.ProgressBar
 	msgPb       *pb.ProgressBar
 	subLevelPbs []*pb.ProgressBar
@@ -127,14 +135,18 @@ type terminalProgressBar struct {
 
 // NewTerminalProgressBar creates a new default pb3 based progressbar suitable for
 // most terminals.
-func NewTerminalProgressBar() (ProgressBar, error) {
+func NewTerminalProgressBar(config *ProgressConfig) (ProgressBar, error) {
 	b := &terminalProgressBar{
-		out: osStderr(),
+		bytes: config.Bytes,
+		speed: config.Speed,
+		out:   osStderr(),
 	}
 	b.spinnerPb = pb.New(0)
 	b.spinnerPb.SetTemplate(`[{{ (cycle . "|" "/" "-" "\\") }}] {{ string . "spinnerMsg" }}`)
-	b.msgPb = pb.New(0)
-	b.msgPb.SetTemplate(`Message: {{ string . "msg" }}`)
+	if config.WithMsg {
+		b.msgPb = pb.New(0)
+		b.msgPb.SetTemplate(`Message: {{ string . "msg" }}`)
+	}
 	return b, nil
 }
 
@@ -148,6 +160,9 @@ func (b *terminalProgressBar) SetProgress(subLevel int, msg string, done int, to
 	case subLevel == len(b.subLevelPbs):
 		apb := pb.New(0)
 		progressBarTmpl := `[{{ counters . }}] {{ string . "prefix" }} {{ bar .}} {{ percent . }}`
+		if b.speed {
+			progressBarTmpl += ` {{ speed . }}`
+		}
 		apb.SetTemplateString(progressBarTmpl)
 		if err := apb.Err(); err != nil {
 			return fmt.Errorf("error setting the progressbarTemplat: %w", err)
@@ -165,6 +180,7 @@ func (b *terminalProgressBar) SetProgress(subLevel int, msg string, done int, to
 	apb.SetTotal(int64(total) + 1)
 	apb.SetCurrent(int64(done) + 1)
 	apb.Set("prefix", msg)
+	apb.Set(pb.Bytes, b.bytes)
 	return nil
 }
 
@@ -173,7 +189,9 @@ func (b *terminalProgressBar) SetPulseMsgf(msg string, args ...any) {
 }
 
 func (b *terminalProgressBar) SetMessagef(msg string, args ...any) {
-	b.msgPb.Set("msg", fmt.Sprintf(msg, args...))
+	if b.msgPb != nil {
+		b.msgPb.Set("msg", fmt.Sprintf(msg, args...))
+	}
 }
 
 func shortenString(msg string) string {
@@ -195,8 +213,10 @@ func (b *terminalProgressBar) render() {
 		fmt.Fprintf(b.out, "%s%s\n", ERASE_LINE, prog.String())
 		renderedLines++
 	}
-	fmt.Fprintf(b.out, "%s%s\n", ERASE_LINE, shortenString(b.msgPb.String()))
-	renderedLines++
+	if b.msgPb != nil {
+		fmt.Fprintf(b.out, "%s%s\n", ERASE_LINE, shortenString(b.msgPb.String()))
+		renderedLines++
+	}
 	fmt.Fprint(b.out, cursorUp(renderedLines))
 }
 
@@ -238,8 +258,10 @@ func (b *terminalProgressBar) Err() error {
 	if err := b.spinnerPb.Err(); err != nil {
 		errs = append(errs, fmt.Errorf("error on spinner progressbar: %w", err))
 	}
-	if err := b.msgPb.Err(); err != nil {
-		errs = append(errs, fmt.Errorf("error on spinner progressbar: %w", err))
+	if b.msgPb != nil {
+		if err := b.msgPb.Err(); err != nil {
+			errs = append(errs, fmt.Errorf("error on spinner progressbar: %w", err))
+		}
 	}
 	for _, pb := range b.subLevelPbs {
 		if err := pb.Err(); err != nil {
@@ -389,8 +411,8 @@ type fileProgressItem struct {
 // NewFileProgressBar starts a new "file" progressbar that will write any progress to a file.
 // Messages without progress information are ignored. The file progress never writes out subprogress
 // without the progress above it.
-func NewFileProgressBar(path string) (ProgressBar, error) {
-	b := &fileProgressBar{path: path}
+func NewFileProgressBar(config *ProgressConfig) (ProgressBar, error) {
+	b := &fileProgressBar{path: config.FilePath}
 	return b, nil
 }
 

pkg/progress/progress_test.go

@@ -109,22 +109,63 @@ func TestTermProgress(t *testing.T) {
 	restore := progress.MockOsStderr(&buf)
 	defer restore()
 
-	pbar, err := progress.NewTerminalProgressBar()
-	assert.NoError(t, err)
+	for _, tc := range []struct {
+		name     string
+		conf     progress.ProgressConfig
+		expected []string
+	}{
+		{
+			name: "default options",
+			conf: progress.ProgressConfig{},
+			expected: []string{
+				"[1 / 6] set-progress-msg",
+				"[|] pulse-msg\n",
+			},
+		},
+		{
+			name: "with message",
+			conf: progress.ProgressConfig{
+				WithMsg: true,
+			},
+			expected: []string{
+				"[1 / 6] set-progress-msg",
+				"[|] pulse-msg\n",
+				"Message: some-message\n",
+			},
+		},
+		{
+			name: "with speed & bytes",
+			conf: progress.ProgressConfig{
+				Bytes: true,
+				Speed: true,
+			},
+			expected: []string{
+				"[1 B / 6 B] set-progress-msg",
+				"p/s\n",
+				"[|] pulse-msg\n",
+			},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			buf.Reset()
 
-	pbar.Start()
-	pbar.SetPulseMsgf("pulse-msg")
-	pbar.SetMessagef("some-message")
-	err = pbar.SetProgress(0, "set-progress-msg", 0, 5)
-	assert.NoError(t, err)
-	pbar.Stop()
-	assert.NoError(t, pbar.(*progress.TerminalProgressBar).Err())
+			pbar, err := progress.NewTerminalProgressBar(&tc.conf)
+			assert.NoError(t, err)
 
-	assert.Contains(t, buf.String(), "[1 / 6] set-progress-msg")
-	assert.Contains(t, buf.String(), "[|] pulse-msg\n")
-	assert.Contains(t, buf.String(), "Message: some-message\n")
-	// check shutdown
-	assert.Contains(t, buf.String(), progress.CURSOR_SHOW)
+			pbar.Start()
+			pbar.SetPulseMsgf("pulse-msg")
+			pbar.SetMessagef("some-message")
+			err = pbar.SetProgress(0, "set-progress-msg", 0, 5)
+			assert.NoError(t, err)
+			pbar.Stop()
+			assert.NoError(t, pbar.(*progress.TerminalProgressBar).Err())
+			for _, exp := range tc.expected {
+				assert.Contains(t, buf.String(), exp)
+			}
+			// check shutdown
+			assert.Contains(t, buf.String(), progress.CURSOR_SHOW)
+		})
+	}
 }
 
 func TestProgressNewAutoselect(t *testing.T) {
@@ -160,7 +201,9 @@ func TestFileProgress(t *testing.T) {
 	testDir := t.TempDir()
 	progressFile := filepath.Join(testDir, "progress")
 
-	pbar, err := progress.NewFileProgressBar(progressFile)
+	pbar, err := progress.NewFileProgressBar(&progress.ProgressConfig{
+		FilePath: progressFile,
+	})
 	assert.NoError(t, err)
 	pbar.Start()
 	_, err = os.Stat(progressFile)