Codebase list golang-github-vbauerster-mpb / 9436aa1
proxyWriterTo Vladimir Bauer 6 years ago
3 changed file(s) with 74 addition(s) and 16 deletion(s). Raw diff Collapse all Expand all
+5
-1
bar.go less more
140140 if !ok {
141141 rc = ioutil.NopCloser(r)
142142 }
143 return &proxyReader{rc, b, time.Now()}
143 prox := &proxyReader{rc, b, time.Now()}
144 if wt, ok := r.(io.WriterTo); ok {
145 return &proxyWriterTo{prox, wt}
146 }
147 return prox
144148 }
145149
146150 // ID returs id of the bar.
+24
-7
proxyreader.go less more
44 "time"
55 )
66
7 // proxyReader is io.Reader wrapper, for proxy read bytes
87 type proxyReader struct {
98 io.ReadCloser
109 bar *Bar
1110 iT time.Time
1211 }
1312
14 func (pr *proxyReader) Read(p []byte) (n int, err error) {
15 n, err = pr.ReadCloser.Read(p)
13 func (prox *proxyReader) Read(p []byte) (n int, err error) {
14 n, err = prox.ReadCloser.Read(p)
1615 if n > 0 {
17 pr.bar.IncrBy(n, time.Since(pr.iT))
18 pr.iT = time.Now()
16 prox.bar.IncrBy(n, time.Since(prox.iT))
17 prox.iT = time.Now()
1918 }
2019 if err == io.EOF {
2120 go func() {
22 current := pr.bar.Current()
23 pr.bar.SetTotal(current, true)
21 prox.bar.SetTotal(prox.bar.Current(), true)
2422 }()
2523 }
2624 return
2725 }
26
27 type proxyWriterTo struct {
28 *proxyReader
29 wt io.WriterTo
30 }
31
32 func (prox *proxyWriterTo) WriteTo(w io.Writer) (n int64, err error) {
33 n, err = prox.wt.WriteTo(w)
34 if n > 0 {
35 prox.bar.IncrInt64(n, time.Since(prox.iT))
36 prox.iT = time.Now()
37 }
38 if err == io.EOF {
39 go func() {
40 prox.bar.SetTotal(prox.bar.Current(), true)
41 }()
42 }
43 return
44 }
+45
-8
proxyreader_test.go less more
00 package mpb_test
11
22 import (
3 "bytes"
34 "io"
45 "io/ioutil"
56 "strings"
2728 }
2829
2930 func TestProxyReader(t *testing.T) {
30
3131 p := mpb.New(mpb.WithOutput(ioutil.Discard))
3232
33 reader := &testReader{Reader: strings.NewReader(content)}
33 tReader := &testReader{strings.NewReader(content), false}
3434
35 total := len(content)
36 bar := p.AddBar(100, mpb.TrimSpace())
35 bar := p.AddBar(int64(len(content)), mpb.TrimSpace())
3736
38 written, err := io.Copy(ioutil.Discard, bar.ProxyReader(reader))
37 var buf bytes.Buffer
38 _, err := io.Copy(&buf, bar.ProxyReader(tReader))
3939 if err != nil {
4040 t.Errorf("Error copying from reader: %+v\n", err)
4141 }
4242
4343 p.Wait()
4444
45 if !reader.called {
45 if !tReader.called {
4646 t.Error("Read not called")
4747 }
4848
49 if written != int64(total) {
50 t.Errorf("Expected written: %d, got: %d\n", total, written)
49 if got := buf.String(); got != content {
50 t.Errorf("Expected content: %s, got: %s\n", content, got)
5151 }
5252 }
53
54 type testWriterTo struct {
55 io.Reader
56 wt io.WriterTo
57 called bool
58 }
59
60 func (wt *testWriterTo) WriteTo(w io.Writer) (n int64, err error) {
61 wt.called = true
62 return wt.wt.WriteTo(w)
63 }
64
65 func TestProxyWriterTo(t *testing.T) {
66 p := mpb.New(mpb.WithOutput(ioutil.Discard))
67
68 var reader io.Reader = strings.NewReader(content)
69 wt := reader.(io.WriterTo)
70 tReader := &testWriterTo{reader, wt, false}
71
72 bar := p.AddBar(int64(len(content)), mpb.TrimSpace())
73
74 var buf bytes.Buffer
75 _, err := io.Copy(&buf, bar.ProxyReader(tReader))
76 if err != nil {
77 t.Errorf("Error copying from reader: %+v\n", err)
78 }
79
80 p.Wait()
81
82 if !tReader.called {
83 t.Error("WriteTo not called")
84 }
85
86 if got := buf.String(); got != content {
87 t.Errorf("Expected content: %s, got: %s\n", content, got)
88 }
89 }