55 lines
1.1 KiB
Go
55 lines
1.1 KiB
Go
package limiter_test
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
"io/ioutil"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/influxdata/influxdb/v2/pkg/limiter"
|
|
)
|
|
|
|
func TestWriter_Limited(t *testing.T) {
|
|
r := bytes.NewReader(bytes.Repeat([]byte{0}, 1024*1024))
|
|
|
|
limit := 512 * 1024
|
|
w := limiter.NewWriter(nopWriteCloser{ioutil.Discard}, 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)
|
|
}
|
|
}
|
|
|
|
func TestWriter_Limiter_ExceedBurst(t *testing.T) {
|
|
limit := 10
|
|
burstLimit := 20
|
|
|
|
twentyOneBytes := make([]byte, 21)
|
|
|
|
b := nopWriteCloser{bytes.NewBuffer(nil)}
|
|
|
|
w := limiter.NewWriter(b, limit, burstLimit)
|
|
n, err := w.Write(twentyOneBytes)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if n != len(twentyOneBytes) {
|
|
t.Errorf("exected %d bytes written, but got %d", len(twentyOneBytes), n)
|
|
}
|
|
}
|
|
|
|
type nopWriteCloser struct {
|
|
io.Writer
|
|
}
|
|
|
|
func (d nopWriteCloser) Close() error { return nil }
|