1 files changed, 101 insertions, 0 deletions
diff --git a/vendor/gopkg.in/throttled/throttled.v1/rate_test.go b/vendor/gopkg.in/throttled/throttled.v1/rate_test.go
new file mode 100644
index 000000000..67dea74b1
--- /dev/null
+++ b/vendor/gopkg.in/throttled/throttled.v1/rate_test.go
@@ -0,0 +1,101 @@
+package throttled
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"testing"
+	"time"
+)
+
+const deniedStatus = 429
+
+// Simple memory store for tests, unsafe for concurrent access
+type mapStore struct {
+	cnt map[string]int
+	ts  map[string]time.Time
+}
+
+func newMapStore() *mapStore {
+	return &mapStore{
+		make(map[string]int),
+		make(map[string]time.Time),
+	}
+}
+func (ms *mapStore) Incr(key string, window time.Duration) (int, int, error) {
+	if _, ok := ms.cnt[key]; !ok {
+		return 0, 0, ErrNoSuchKey
+	}
+	ms.cnt[key]++
+	ts := ms.ts[key]
+	return ms.cnt[key], RemainingSeconds(ts, window), nil
+}
+func (ms *mapStore) Reset(key string, win time.Duration) error {
+	ms.cnt[key] = 1
+	ms.ts[key] = time.Now().UTC()
+	return nil
+}
+
+func TestRateLimit(t *testing.T) {
+	quota := Q{5, 5 * time.Second}
+	cases := []struct {
+		limit, remain, reset, status int
+	}{
+		0: {5, 4, 5, 200},
+		1: {5, 3, 4, 200},
+		2: {5, 2, 4, 200},
+		3: {5, 1, 3, 200},
+		4: {5, 0, 3, 200},
+		5: {5, 0, 2, deniedStatus},
+	}
+	// Limit the requests to 2 per second
+	th := Interval(PerSec(2), 0, nil, 0)
+	// Rate limit
+	rl := RateLimit(quota, nil, newMapStore())
+	// Create the stats
+	st := &stats{}
+	// Create the handler
+	h := th.Throttle(rl.Throttle(st))
+
+	// Start the server
+	srv := httptest.NewServer(h)
+	defer srv.Close()
+	for i, c := range cases {
+		callRateLimited(t, i, c.limit, c.remain, c.reset, c.status, srv.URL)
+	}
+	// Wait 3 seconds and call again, should start a new window
+	time.Sleep(3 * time.Second)
+	callRateLimited(t, len(cases), 5, 4, 5, 200, srv.URL)
+}
+
+func callRateLimited(t *testing.T, i, limit, remain, reset, status int, url string) {
+	res, err := http.Get(url)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer res.Body.Close()
+	// Assert status code
+	if status != res.StatusCode {
+		t.Errorf("%d: expected status %d, got %d", i, status, res.StatusCode)
+	}
+	// Assert headers
+	if v := res.Header.Get("X-RateLimit-Limit"); v != strconv.Itoa(limit) {
+		t.Errorf("%d: expected limit header to be %d, got %s", i, limit, v)
+	}
+	if v := res.Header.Get("X-RateLimit-Remaining"); v != strconv.Itoa(remain) {
+		t.Errorf("%d: expected remain header to be %d, got %s", i, remain, v)
+	}
+	// Allow 1 second wiggle room
+	v := res.Header.Get("X-RateLimit-Reset")
+	vi, _ := strconv.Atoi(v)
+	if vi < reset-1 || vi > reset+1 {
+		t.Errorf("%d: expected reset header to be close to %d, got %d", i, reset, vi)
+	}
+	if status == deniedStatus {
+		v := res.Header.Get("Retry-After")
+		vi, _ := strconv.Atoi(v)
+		if vi < reset-1 || vi > reset+1 {
+			t.Errorf("%d: expected retry after header to be close to %d, got %d", i, reset, vi)
+		}
+	}
+}