diff --git a/go.mod b/go.mod index 1efa745..6d8333d 100644 --- a/go.mod +++ b/go.mod @@ -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 ) @@ -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 diff --git a/internal/http/ratelimit_test.go b/internal/http/ratelimit_test.go index e620c0b..5630928 100644 --- a/internal/http/ratelimit_test.go +++ b/internal/http/ratelimit_test.go @@ -1,6 +1,7 @@ package http import ( + "fmt" "net/http" "net/http/httptest" "sync" @@ -8,6 +9,7 @@ import ( "github.com/lucasrod16/oss-contribute/internal/cache" "github.com/stretchr/testify/require" + "golang.org/x/sync/errgroup" ) func TestRateLimiter(t *testing.T) { @@ -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")