35 lines
742 B
Go
35 lines
742 B
Go
package limiter_test
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/influxdata/influxdb/pkg/limiter"
|
|
)
|
|
|
|
func TestWriter_Limited(t *testing.T) {
|
|
r := bytes.NewReader(bytes.Repeat([]byte{0}, 1024*1024))
|
|
|
|
limit := 512 * 1024
|
|
w := limiter.NewWriter(discardCloser{}, limit, 10*1024*1024)
|
|
|
|
start := time.Now()
|
|
n, err := io.Copy(w, r)
|
|
elapsed := time.Since(start)
|
|
if err != nil {
|
|
t.Error("copy error: ", err)
|
|
}
|
|
|
|
rate := float64(n) / elapsed.Seconds()
|
|
if rate > float64(limit) {
|
|
t.Errorf("rate limit mismath: exp %f, got %f", float64(limit), rate)
|
|
}
|
|
}
|
|
|
|
type discardCloser struct{}
|
|
|
|
func (d discardCloser) Write(b []byte) (int, error) { return len(b), nil }
|
|
func (d discardCloser) Close() error { return nil }
|