File: http_test.go

package info (click to toggle)
golang-github-throttled-throttled 2.2.4-1
  • links: PTS, VCS
  • area: main
  • in suites: bullseye, sid
  • size: 172 kB
  • sloc: makefile: 23
file content (112 lines) | stat: -rw-r--r-- 2,715 bytes parent folder | download | duplicates (2)
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
104
105
106
107
108
109
110
111
112
package throttled_test

import (
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/throttled/throttled"
)

type stubLimiter struct {
}

func (sl *stubLimiter) RateLimit(key string, quantity int) (bool, throttled.RateLimitResult, error) {
	switch key {
	case "limit":
		result := throttled.RateLimitResult{
			Limit:      -1,
			Remaining:  -1,
			ResetAfter: -1,
			RetryAfter: time.Minute,
		}
		return true, result, nil
	case "error":
		result := throttled.RateLimitResult{}
		return false, result, errors.New("stubLimiter error")
	default:
		result := throttled.RateLimitResult{
			Limit:      1,
			Remaining:  2,
			ResetAfter: time.Minute,
			RetryAfter: -1,
		}
		return false, result, nil
	}
}

type pathGetter struct{}

func (*pathGetter) Key(r *http.Request) string {
	return r.URL.Path
}

type httpTestCase struct {
	path    string
	code    int
	headers map[string]string
}

func TestHTTPRateLimiter(t *testing.T) {
	limiter := throttled.HTTPRateLimiter{
		RateLimiter: &stubLimiter{},
		VaryBy:      &pathGetter{},
	}

	handler := limiter.RateLimit(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
	}))

	runHTTPTestCases(t, handler, []httpTestCase{
		{"ok", 200, map[string]string{"X-Ratelimit-Limit": "1", "X-Ratelimit-Remaining": "2", "X-Ratelimit-Reset": "60"}},
		{"error", 500, map[string]string{}},
		{"limit", 429, map[string]string{"Retry-After": "60"}},
	})
}

func TestCustomHTTPRateLimiterHandlers(t *testing.T) {
	limiter := throttled.HTTPRateLimiter{
		RateLimiter: &stubLimiter{},
		VaryBy:      &pathGetter{},
		DeniedHandler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			http.Error(w, "custom limit exceeded", 400)
		}),
		Error: func(w http.ResponseWriter, r *http.Request, err error) {
			http.Error(w, "custom internal error", 501)
		},
	}

	handler := limiter.RateLimit(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
	}))

	runHTTPTestCases(t, handler, []httpTestCase{
		{"limit", 400, map[string]string{}},
		{"error", 501, map[string]string{}},
	})
}

func runHTTPTestCases(t *testing.T, h http.Handler, cs []httpTestCase) {
	for i, c := range cs {
		req, err := http.NewRequest("GET", c.path, nil)
		if err != nil {
			t.Fatal(err)
		}

		rr := httptest.NewRecorder()
		h.ServeHTTP(rr, req)
		if have, want := rr.Code, c.code; have != want {
			t.Errorf("Expected request %d at %s to return %d but got %d",
				i, c.path, want, have)
		}

		for name, want := range c.headers {
			if have := rr.HeaderMap.Get(name); have != want {
				t.Errorf("Expected request %d at %s to have header '%s: %s' but got '%s'",
					i, c.path, name, want, have)
			}
		}
	}
}