diff --git a/progressreader/progressreader.go b/progressreader/progressreader.go index 652831b..908e8ee 100644 --- a/progressreader/progressreader.go +++ b/progressreader/progressreader.go @@ -23,6 +23,7 @@ type Config struct { func New(newReader Config) *Config { return &newReader } + func (config *Config) Read(p []byte) (n int, err error) { read, err := config.In.Read(p) config.Current += read @@ -34,17 +35,30 @@ func (config *Config) Read(p []byte) (n int, err error) { } } if config.Current-config.LastUpdate > updateEvery || err != nil { - config.Out.Write(config.Formatter.FormatProgress(config.ID, config.Action, &jsonmessage.JSONProgress{Current: config.Current, Total: config.Size})) + updateProgress(config) config.LastUpdate = config.Current } - // Send newline when complete - if config.NewLines && err != nil && read == 0 { - config.Out.Write(config.Formatter.FormatStatus("", "")) + + if err != nil && read == 0 { + updateProgress(config) + if config.NewLines { + config.Out.Write(config.Formatter.FormatStatus("", "")) + } } return read, err } + func (config *Config) Close() error { - config.Current = config.Size - config.Out.Write(config.Formatter.FormatProgress(config.ID, config.Action, &jsonmessage.JSONProgress{Current: config.Current, Total: config.Size})) + if config.Current < config.Size { + //print a full progress bar when closing prematurely + config.Current = config.Size + updateProgress(config) + } return config.In.Close() } + +func updateProgress(config *Config) { + progress := jsonmessage.JSONProgress{Current: config.Current, Total: config.Size} + fmtMessage := config.Formatter.FormatProgress(config.ID, config.Action, &progress) + config.Out.Write(fmtMessage) +} diff --git a/progressreader/progressreader_test.go b/progressreader/progressreader_test.go new file mode 100644 index 0000000..fdf40cb --- /dev/null +++ b/progressreader/progressreader_test.go @@ -0,0 +1,94 @@ +package progressreader + +import ( + "bufio" + "bytes" + "io" + "io/ioutil" + "testing" + + "github.com/docker/docker/pkg/streamformatter" +) + +func TestOutputOnPrematureClose(t *testing.T) { + var outBuf bytes.Buffer + content := []byte("TESTING") + reader := ioutil.NopCloser(bytes.NewReader(content)) + writer := bufio.NewWriter(&outBuf) + + prCfg := Config{ + In: reader, + Out: writer, + Formatter: streamformatter.NewStreamFormatter(), + Size: len(content), + NewLines: true, + ID: "Test", + Action: "Read", + } + pr := New(prCfg) + + part := make([]byte, 4, 4) + _, err := io.ReadFull(pr, part) + if err != nil { + pr.Close() + t.Fatal(err) + } + + if err := writer.Flush(); err != nil { + pr.Close() + t.Fatal(err) + } + + tlen := outBuf.Len() + pr.Close() + if err := writer.Flush(); err != nil { + t.Fatal(err) + } + + if outBuf.Len() == tlen { + t.Fatalf("Expected some output when closing prematurely") + } +} + +func TestCompleteSilently(t *testing.T) { + var outBuf bytes.Buffer + content := []byte("TESTING") + reader := ioutil.NopCloser(bytes.NewReader(content)) + writer := bufio.NewWriter(&outBuf) + + prCfg := Config{ + In: reader, + Out: writer, + Formatter: streamformatter.NewStreamFormatter(), + Size: len(content), + NewLines: true, + ID: "Test", + Action: "Read", + } + pr := New(prCfg) + + out, err := ioutil.ReadAll(pr) + if err != nil { + pr.Close() + t.Fatal(err) + } + if string(out) != "TESTING" { + pr.Close() + t.Fatalf("Unexpected output %q from reader", string(out)) + } + + if err := writer.Flush(); err != nil { + pr.Close() + t.Fatal(err) + } + + tlen := outBuf.Len() + pr.Close() + if err := writer.Flush(); err != nil { + t.Fatal(err) + } + + if outBuf.Len() > tlen { + t.Fatalf("Should have closed silently when read is complete") + } +}