fix/608 - Fix flaky Throttle middleware test by synchronizing token usage (#1016)

* chore: remove comment from TestThrottleRetryAfter test * test: fix flaky TestThrottleRetryAfter by synchronizing token usage

Author: OtavioBernardes

middleware/throttle_test.go

@@ -199,55 +199,67 @@ func TestThrottleMaximum(t *testing.T) {
 wg.Wait()
 }
 
-// NOTE: test is disabled as it requires some refactoring. It is prone to intermittent failure.
-/*func TestThrottleRetryAfter(t *testing.T) {
+func TestThrottleRetryAfter(t *testing.T) {
 r := chi.NewRouter()
+retryAfterFn := func(ctxDone bool) time.Duration { return time.Hour }
 
-retryAfterFn := func(ctxDone bool) time.Duration { return time.Hour * 1 }
-r.Use(ThrottleWithOpts(ThrottleOpts{Limit: 10, RetryAfterFn: retryAfterFn}))
+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)
-time.Sleep(time.Second * 4) // Expensive operation.
-w.Write(testContent)
+w.Write([]byte("ok"))
 })
 
 server := httptest.NewServer(r)
 defer server.Close()
+client := http.Client{}
 
-client := http.Client{
-Timeout: time.Second * 60, // Maximum waiting time.
+type result struct {
+status int
+header http.Header
 }
 
 var wg sync.WaitGroup
+totalRequests := 10
+resultsCh := make(chan result, totalRequests)
 
-for i := 0; i < 10; i++ {
+for i := 0; i < totalRequests; i++ {
 wg.Add(1)
-go func(i int) {
+go func() {
 defer wg.Done()
-
-res, err := client.Get(server.URL)
-assertNoError(t, err)
-assertEqual(t, http.StatusOK, res.StatusCode)
-}(i)
+res, _ := client.Get(server.URL)
+resultsCh <- result{status: res.StatusCode, header: res.Header}
+}()
 }
 
-time.Sleep(time.Second * 1)
-
-for i := 0; i < 10; i++ {
-wg.Add(1)
-go func(i int) {
-defer wg.Done()
-
-res, err := client.Get(server.URL)
-assertNoError(t, err)
-assertEqual(t, http.StatusTooManyRequests, res.StatusCode)
-assertEqual(t, res.Header.Get("Retry-After"), "3600")
-}(i)
+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
+}
 }
 
-wg.Wait()
-}*/
+assertEqual(t, 5, count200)
+assertEqual(t, 5, count429)
+}
 
 func TestThrottleCustomStatusCode(t *testing.T) {
 const timeout = time.Second * 3