// We've hit the limit and were able to proceed on the next time run
func (s *LimiterSuite) TestOptions(c *C) {
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
rates := NewRateSet()
rates.Add(time.Second, 1, 1)
errHandler := utils.ErrorHandlerFunc(func(w http.ResponseWriter, req *http.Request, err error) {
w.WriteHeader(http.StatusTeapot)
w.Write([]byte(http.StatusText(http.StatusTeapot)))
})
buf := &bytes.Buffer{}
log := utils.NewFileLogger(buf, utils.INFO)
l, err := New(handler, headerLimit, rates, ErrorHandler(errHandler), Logger(log), Clock(s.clock))
c.Assert(err, IsNil)
srv := httptest.NewServer(l)
defer srv.Close()
re, _, err := testutils.Get(srv.URL, testutils.Header("Source", "a"))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
re, _, err = testutils.Get(srv.URL, testutils.Header("Source", "a"))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusTeapot)
c.Assert(len(buf.String()), Not(Equals), 0)
}
// We've hit the limit and were able to proceed on the next time run
func (s *LimiterSuite) TestHitLimit(c *C) {
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
rates := NewRateSet()
rates.Add(time.Second, 1, 1)
l, err := New(handler, headerLimit, rates, Clock(s.clock))
c.Assert(err, IsNil)
srv := httptest.NewServer(l)
defer srv.Close()
re, _, err := testutils.Get(srv.URL, testutils.Header("Source", "a"))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
// Next request from the same source hits rate limit
re, _, err = testutils.Get(srv.URL, testutils.Header("Source", "a"))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, 429)
// Second later, the request from this ip will succeed
s.clock.Sleep(time.Second)
re, _, err = testutils.Get(srv.URL, testutils.Header("Source", "a"))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
}
// If configMapper returns empty rates, then the default rate is applied.
func (s *LimiterSuite) TestExtractorEmpty(c *C) {
// Given
extractor := func(*http.Request) (*RateSet, error) {
return NewRateSet(), nil
}
rates := NewRateSet()
rates.Add(time.Second, 1, 1)
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
l, err := New(handler, headerLimit, rates, Clock(s.clock), ExtractRates(RateExtractorFunc(extractor)))
c.Assert(err, IsNil)
srv := httptest.NewServer(l)
defer srv.Close()
// When/Then: The default rate is applied, which 1 req/second
re, _, err := testutils.Get(srv.URL, testutils.Header("Source", "a"))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
re, _, err = testutils.Get(srv.URL, testutils.Header("Source", "a"))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, 429)
s.clock.Sleep(time.Second)
re, _, err = testutils.Get(srv.URL, testutils.Header("Source", "a"))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
}
// We've failed to extract client ip
func (s *LimiterSuite) TestFailure(c *C) {
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
rates := NewRateSet()
rates.Add(time.Second, 1, 1)
l, err := New(handler, faultyExtract, rates, Clock(s.clock))
c.Assert(err, IsNil)
srv := httptest.NewServer(l)
defer srv.Close()
re, _, err := testutils.Get(srv.URL, testutils.Header("Source", "a"))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusInternalServerError)
}