diff --git a/cmd/image-builder/main.go b/cmd/image-builder/main.go index a3da3dcbc8..8ea90aadd6 100644 --- a/cmd/image-builder/main.go +++ b/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))) diff --git a/cmd/image-builder/upload.go b/cmd/image-builder/upload.go index da3ff99328..c82b61e9a2 100644 --- a/cmd/image-builder/upload.go +++ b/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 } diff --git a/cmd/image-builder/upload_test.go b/cmd/image-builder/upload_test.go index a2a5124a41..4941bc3363 100644 --- a/cmd/image-builder/upload_test.go +++ b/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") } } diff --git a/pkg/progress/progress.go b/pkg/progress/progress.go index 217fca9a4d..ea77cd592e 100644 --- a/pkg/progress/progress.go +++ b/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 } diff --git a/pkg/progress/progress_test.go b/pkg/progress/progress_test.go index c45278bd49..555373d773 100644 --- a/pkg/progress/progress_test.go +++ b/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) diff --git a/pkg/progress/proxy_reader.go b/pkg/progress/proxy_reader.go new file mode 100644 index 0000000000..a021261ccf --- /dev/null +++ b/pkg/progress/proxy_reader.go @@ -0,0 +1,42 @@ +package progress + +import ( + "io" +) + +type Reader struct { + io.Reader + pb ProgressBar + done int + total int +} + +func NewProxyReader(r io.Reader, total int, pb ProgressBar) (*Reader, error) { + err := pb.SetProgress(0, "Uploading", 0, total) + if err != nil { + return nil, err + } + return &Reader{ + Reader: r, + pb: pb, + done: 0, + total: total, + }, nil +} + +func (r *Reader) Read(p []byte) (int, error) { + n, err := r.Reader.Read(p) + r.done += n + if err != nil { + return n, err + } + err = r.pb.SetProgress(0, "", r.done, r.total) + return n, err +} + +func (r *Reader) Close() error { + if closer, ok := r.Reader.(io.Closer); ok { + return closer.Close() + } + return nil +} diff --git a/pkg/progress/proxy_reader_test.go b/pkg/progress/proxy_reader_test.go new file mode 100644 index 0000000000..f7baa3017f --- /dev/null +++ b/pkg/progress/proxy_reader_test.go @@ -0,0 +1,37 @@ +package progress_test + +import ( + "bytes" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/osbuild/image-builder/pkg/progress" +) + +func TestProxyReader(t *testing.T) { + var progressBuf bytes.Buffer + restore := progress.MockOsStderr(&progressBuf) + defer restore() + + pbar, err := progress.NewDebugProgressBar() + assert.NoError(t, err) + + var readBuf bytes.Buffer + size, err := readBuf.Write([]byte("duck")) + assert.NoError(t, err) + assert.Equal(t, 4, size) + proxyReader, err := progress.NewProxyReader(bytes.NewReader(readBuf.Bytes()), size, pbar) + assert.NoError(t, err) + + out := make([]byte, 256) + nRead, err := proxyReader.Read(out) + assert.NoError(t, err) + assert.Equal(t, 4, nRead) + assert.Equal(t, "duck", string(out[:nRead])) + + lines := strings.Split(progressBuf.String(), "\n") + assert.Equal(t, "[0 / 4] Uploading", lines[0]) + assert.Equal(t, "[4 / 4] ", lines[1]) +}