Skip to content

Commit

Permalink
Use switch statement and errgroup in rate limit test
Browse files Browse the repository at this point in the history
Signed-off-by: Lucas Rodriguez <[email protected]>
  • Loading branch information
lucasrod16 committed Sep 29, 2024
1 parent 7dde0f9 commit c4be9ab
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 13 deletions.
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ require (
cloud.google.com/go/storage v1.43.0
github.com/google/go-github/v64 v64.0.0
github.com/stretchr/testify v1.9.0
golang.org/x/sync v0.7.0
golang.org/x/time v0.6.0
)

Expand Down Expand Up @@ -36,7 +37,6 @@ require (
golang.org/x/crypto v0.24.0 // indirect
golang.org/x/net v0.26.0 // indirect
golang.org/x/oauth2 v0.21.0 // indirect
golang.org/x/sync v0.7.0 // indirect
golang.org/x/sys v0.21.0 // indirect
golang.org/x/text v0.16.0 // indirect
google.golang.org/api v0.187.0 // indirect
Expand Down
30 changes: 18 additions & 12 deletions internal/http/ratelimit_test.go
Original file line number Diff line number Diff line change
@@ -1,13 +1,15 @@
package http

import (
"fmt"
"net/http"
"net/http/httptest"
"sync"
"testing"

"github.com/lucasrod16/oss-contribute/internal/cache"
"github.com/stretchr/testify/require"
"golang.org/x/sync/errgroup"
)

func TestRateLimiter(t *testing.T) {
Expand All @@ -18,31 +20,35 @@ func TestRateLimiter(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "/repos", nil)
req.Header.Set("X-Forwarded-For", "192.168.1.1")

var wg sync.WaitGroup
var g errgroup.Group
var mu sync.Mutex

successCount := 0
failCount := 0

// send 20 requests concurrently to trigger the rate limit
for i := 0; i < 20; i++ {
wg.Add(1)
go func() {
defer wg.Done()
g.Go(func() error {
rr := httptest.NewRecorder()
rl.Limit(GetRepos(c)).ServeHTTP(rr, req)

if rr.Code == http.StatusOK {
mu.Lock()
mu.Lock()
defer mu.Unlock()

switch rr.Code {
case http.StatusOK:
successCount++
mu.Unlock()
} else if rr.Code == http.StatusTooManyRequests {
mu.Lock()
case http.StatusTooManyRequests:
failCount++
mu.Unlock()
default:
return fmt.Errorf("unexpected status code: %d", rr.Code)
}
}()
return nil
})
}
if err := g.Wait(); err != nil {
t.Fatal(err)
}
wg.Wait()

require.Equal(t, 10, successCount, "Expected 10 successful requests")
require.Equal(t, 10, failCount, "Expected 10 failed requests")
Expand Down

0 comments on commit c4be9ab

Please sign in to comment.