| 3 | 3 |
"bytes"
|
| 4 | 4 |
"io"
|
| 5 | 5 |
"io/ioutil"
|
| 6 | |
"net/http"
|
| 7 | |
"net/http/httptest"
|
| 8 | 6 |
"strings"
|
| 9 | 7 |
"testing"
|
| 10 | 8 |
|
|
| 29 | 27 |
bar := p.AddBar(100, mpb.BarTrim())
|
| 30 | 28 |
preader := bar.ProxyReader(reader)
|
| 31 | 29 |
|
|
30 |
if _, ok := preader.(io.Closer); !ok {
|
|
31 |
t.Error("type assertion to io.Closer is not ok")
|
|
32 |
}
|
|
33 |
|
| 32 | 34 |
written, err := io.Copy(ioutil.Discard, preader)
|
| 33 | 35 |
if err != nil {
|
| 34 | 36 |
t.Errorf("Error copying from reader: %+v\n", err)
|
|
| 39 | 41 |
if written != int64(total) {
|
| 40 | 42 |
t.Errorf("Expected written: %d, got: %d\n", total, written)
|
| 41 | 43 |
}
|
| 42 | |
|
| 43 | |
// underlying reader is not Closer
|
| 44 | |
err = preader.Close()
|
| 45 | |
if err != nil {
|
| 46 | |
t.Errorf("Expected nil error, got: %+v\n", err)
|
| 47 | |
}
|
| 48 | 44 |
}
|
| 49 | |
|
| 50 | |
func TestProxyReaderCloser(t *testing.T) {
|
| 51 | |
var buf bytes.Buffer
|
| 52 | |
p := mpb.New(mpb.WithOutput(&buf))
|
| 53 | |
|
| 54 | |
ts := setupTestHttpServer(content)
|
| 55 | |
defer ts.Close()
|
| 56 | |
|
| 57 | |
url := ts.URL + "/test"
|
| 58 | |
resp, err := http.Get(url)
|
| 59 | |
if err != nil {
|
| 60 | |
t.Errorf("Test server get failure: %s\n", url)
|
| 61 | |
}
|
| 62 | |
|
| 63 | |
total := resp.ContentLength
|
| 64 | |
bar := p.AddBar(total, mpb.BarTrim())
|
| 65 | |
reader := bar.ProxyReader(resp.Body)
|
| 66 | |
|
| 67 | |
// calling reader.Close() will call resp.Body.Close() implicitly
|
| 68 | |
err = reader.Close()
|
| 69 | |
if err != nil {
|
| 70 | |
t.Logf("Error closing resp.Body over reader.Close: %+v\n", err)
|
| 71 | |
t.FailNow()
|
| 72 | |
}
|
| 73 | |
|
| 74 | |
// reading from closed resp.Body
|
| 75 | |
_, err = io.Copy(ioutil.Discard, reader)
|
| 76 | |
if err == nil {
|
| 77 | |
t.Error("Expected read on closed response body error!")
|
| 78 | |
}
|
| 79 | |
}
|
| 80 | |
|
| 81 | |
func setupTestHttpServer(content string) *httptest.Server {
|
| 82 | |
mux := http.NewServeMux()
|
| 83 | |
mux.HandleFunc("/test", func(w http.ResponseWriter, r *http.Request) {
|
| 84 | |
io.WriteString(w, content)
|
| 85 | |
})
|
| 86 | |
return httptest.NewServer(mux)
|
| 87 | |
}
|