forked from wgroeneveld/go-jamming
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
58 lines
1.4 KiB
58 lines
1.4 KiB
package app
|
|
|
|
import (
|
|
"fmt"
|
|
"github.com/gorilla/mux"
|
|
"github.com/stretchr/testify/assert"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
)
|
|
|
|
func TestHitsRateLimitAfterSlammingRequests(t *testing.T) {
|
|
r := mux.NewRouter()
|
|
r.HandleFunc("/endpoint", testFn).Methods("GET")
|
|
r.Use(NewRateLimiter(5, 10).Middleware)
|
|
ts := httptest.NewServer(r)
|
|
|
|
t.Cleanup(ts.Close)
|
|
statusCodes := []int{}
|
|
|
|
for i := 0; i <= 10; i++ {
|
|
client := &http.Client{}
|
|
req, _ := http.NewRequest("GET", fmt.Sprintf("%s/endpoint", ts.URL), nil)
|
|
|
|
resp, err := client.Do(req)
|
|
assert.NoError(t, err)
|
|
statusCodes = append(statusCodes, resp.StatusCode)
|
|
}
|
|
assert.Contains(t, statusCodes, 429)
|
|
}
|
|
|
|
func TestDoesNotHitRateLimitOfSecondEndpointAfterSlammingFirst(t *testing.T) {
|
|
r := mux.NewRouter()
|
|
r.HandleFunc("/endpoint1", testFn).Methods("GET")
|
|
r.HandleFunc("/endpoint2", testFn).Methods("GET")
|
|
r.Use(NewRateLimiter(5, 10).Middleware)
|
|
ts := httptest.NewServer(r)
|
|
|
|
t.Cleanup(ts.Close)
|
|
|
|
for i := 0; i <= 10; i++ {
|
|
client := &http.Client{}
|
|
req, _ := http.NewRequest("GET", fmt.Sprintf("%s/endpoint1", ts.URL), nil)
|
|
client.Do(req)
|
|
}
|
|
for i := 0; i <= 5; i++ {
|
|
client := &http.Client{}
|
|
req, _ := http.NewRequest("GET", fmt.Sprintf("%s/endpoint2", ts.URL), nil)
|
|
resp, err := client.Do(req)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 200, resp.StatusCode)
|
|
}
|
|
}
|
|
|
|
func testFn(w http.ResponseWriter, r *http.Request) {
|
|
w.Write([]byte("ok"))
|
|
}
|