Function TestThrottleRetryAfter

middleware/throttle_test.go:202–262

(t *testing.T)

}

func TestThrottleRetryAfter(t *testing.T) {
	r := chi.NewRouter()
	retryAfterFn := func(ctxDone bool) time.Duration { return time.Hour }

	r.Use(ThrottleWithOpts(ThrottleOpts{
		Limit:        5,
		BacklogLimit: 0,
		RetryAfterFn: retryAfterFn,
	}))

	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(time.Second * 1) // Expensive operation.
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("ok"))
	})

	server := httptest.NewServer(r)
	defer server.Close()
	client := http.Client{}

	type result struct {
		status int
		header http.Header
	}

	var wg sync.WaitGroup
	totalRequests := 10
	resultsCh := make(chan result, totalRequests)

	for i := 0; i < totalRequests; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			res, _ := client.Get(server.URL)
			resultsCh <- result{status: res.StatusCode, header: res.Header}
		}()
	}

	wg.Wait()
	close(resultsCh)

	count200 := 0
	count429 := 0
	for res := range resultsCh {
		switch res.status {
		case http.StatusOK:
			count200++
			continue
		case http.StatusTooManyRequests:
			count429++
			assertEqual(t, "3600", res.header.Get("Retry-After"))
			continue
		default:
			t.Fatalf("Unexpected status code: %d", res.status)
			continue
		}
	}