-
-
Notifications
You must be signed in to change notification settings - Fork 18
/
shapeio_test.go
103 lines (94 loc) · 2.39 KB
/
shapeio_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
package shapeio_test
import (
"bytes"
"io"
"io/ioutil"
"net/http"
"os"
"testing"
"time"
"github.com/dustin/go-humanize"
"github.com/fujiwara/shapeio"
)
var rates = []float64{
500 * 1024, // 500KB/sec
1024 * 1024, // 1MB/sec
10 * 1024 * 1024, // 10MB/sec
50 * 1024 * 1024, // 50MB/sec
}
var srcs = []*bytes.Reader{
bytes.NewReader(bytes.Repeat([]byte{0}, 64*1024)), // 64KB
bytes.NewReader(bytes.Repeat([]byte{1}, 256*1024)), // 256KB
bytes.NewReader(bytes.Repeat([]byte{2}, 1024*1024)), // 1MB
}
func ExampleReader() {
// example for downloading http body with rate limit.
resp, _ := http.Get("http://example.com")
defer resp.Body.Close()
reader := shapeio.NewReader(resp.Body)
reader.SetRateLimit(1024 * 10) // 10KB/sec
io.Copy(ioutil.Discard, reader)
}
func ExampleWriter() {
// example for writing file with rate limit.
src := bytes.NewReader(bytes.Repeat([]byte{0}, 32*1024)) // 32KB
f, _ := os.Create("/tmp/foo")
writer := shapeio.NewWriter(f)
writer.SetRateLimit(1024 * 10) // 10KB/sec
io.Copy(writer, src)
f.Close()
}
func TestRead(t *testing.T) {
for _, src := range srcs {
for _, limit := range rates {
src.Seek(0, 0)
sio := shapeio.NewReader(src)
sio.SetRateLimit(limit)
start := time.Now()
n, err := io.Copy(ioutil.Discard, sio)
elapsed := time.Since(start)
if err != nil {
t.Error("io.Copy failed", err)
}
realRate := float64(n) / elapsed.Seconds()
if realRate > limit {
t.Errorf("Limit %f but real rate %f", limit, realRate)
}
t.Logf(
"read %s / %s: Real %s/sec Limit %s/sec. (%f %%)",
humanize.IBytes(uint64(n)),
elapsed,
humanize.IBytes(uint64(realRate)),
humanize.IBytes(uint64(limit)),
realRate/limit*100,
)
}
}
}
func TestWrite(t *testing.T) {
for _, src := range srcs {
for _, limit := range rates {
src.Seek(0, 0)
sio := shapeio.NewWriter(ioutil.Discard)
sio.SetRateLimit(limit)
start := time.Now()
n, err := io.Copy(sio, src)
elapsed := time.Since(start)
if err != nil {
t.Error("io.Copy failed", err)
}
realRate := float64(n) / elapsed.Seconds()
if realRate > limit {
t.Errorf("Limit %f but real rate %f", limit, realRate)
}
t.Logf(
"write %s / %s: Real %s/sec Limit %s/sec. (%f %%)",
humanize.IBytes(uint64(n)),
elapsed,
humanize.IBytes(uint64(realRate)),
humanize.IBytes(uint64(limit)),
realRate/limit*100,
)
}
}
}