From ba2cebdfadebcf10d5d27b4f2fa160418e746573 Mon Sep 17 00:00:00 2001 From: naveensrinivasan <172697+naveensrinivasan@users.noreply.github.com> Date: Thu, 29 Jun 2023 20:17:57 -0500 Subject: [PATCH 1/2] :seedling: Improve rate limit handling in roundtripper - Add rate limit testing and handling functionality - Add tests for successful response and Retry-After header set scenarios Signed-off-by: naveensrinivasan <172697+naveensrinivasan@users.noreply.github.com> --- clients/githubrepo/roundtripper/rate_limit.go | 2 +- .../roundtripper/rate_limit_test.go | 91 +++++++++++++++++++ 2 files changed, 92 insertions(+), 1 deletion(-) create mode 100644 clients/githubrepo/roundtripper/rate_limit_test.go diff --git a/clients/githubrepo/roundtripper/rate_limit.go b/clients/githubrepo/roundtripper/rate_limit.go index 950e8f8aa8b..d2ae8378ab1 100644 --- a/clients/githubrepo/roundtripper/rate_limit.go +++ b/clients/githubrepo/roundtripper/rate_limit.go @@ -42,7 +42,7 @@ type rateLimitTransport struct { innerTransport http.RoundTripper } -// Roundtrip handles caching and ratelimiting of responses from GitHub. +// RoundTrip handles caching and rate-limiting of responses from GitHub. func (gh *rateLimitTransport) RoundTrip(r *http.Request) (*http.Response, error) { resp, err := gh.innerTransport.RoundTrip(r) if err != nil { diff --git a/clients/githubrepo/roundtripper/rate_limit_test.go b/clients/githubrepo/roundtripper/rate_limit_test.go new file mode 100644 index 00000000000..d0e17c976cb --- /dev/null +++ b/clients/githubrepo/roundtripper/rate_limit_test.go @@ -0,0 +1,91 @@ +// Copyright 2023 OpenSSF Scorecard Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +package roundtripper + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/ossf/scorecard/v4/log" +) + +func TestRoundTrip(t *testing.T) { + var requestCount int + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Customize the response headers and body based on the test scenario + switch r.URL.Path { + case "/error": + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte("Internal Server Error")) // nolint: errcheck + case "/retry": + requestCount++ + if requestCount == 2 { + // Second request: Return successful response + w.Header().Set("X-RateLimit-Remaining", "10") + w.WriteHeader(http.StatusOK) + w.Write([]byte("Success")) // nolint: errcheck + } else { + // First request: Return Retry-After header + w.Header().Set("Retry-After", "5") + w.WriteHeader(http.StatusTooManyRequests) + w.Write([]byte("Rate Limit Exceeded")) // nolint: errcheck + } + case "/success": + w.Header().Set("X-RateLimit-Remaining", "10") + w.WriteHeader(http.StatusOK) + w.Write([]byte("Success")) // nolint: errcheck + } + })) + defer ts.Close() + + // Create the rateLimitTransport with the test server as the inner transport and a default logger + transport := &rateLimitTransport{ + innerTransport: ts.Client().Transport, + logger: log.NewLogger(log.DefaultLevel), + } + + t.Run("Successful response", func(t *testing.T) { + req, err := http.NewRequest(http.MethodGet, ts.URL+"/success", nil) + if err != nil { + t.Fatalf("Failed to create request: %v", err) + } + + resp, err := transport.RoundTrip(req) + if err != nil { + t.Errorf("Unexpected error: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Errorf("Expected status code %d, got %d", http.StatusOK, resp.StatusCode) + } + }) + + t.Run("Retry-After header set", func(t *testing.T) { + req, err := http.NewRequest(http.MethodGet, ts.URL+"/retry", nil) + if err != nil { + t.Fatalf("Failed to create request: %v", err) + } + + resp, err := transport.RoundTrip(req) + if err != nil { + t.Errorf("Unexpected error: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Errorf("Expected status code %d, got %d", http.StatusOK, resp.StatusCode) + } + if requestCount != 2 { + t.Errorf("Expected 2 requests, got %d", requestCount) + } + }) +} From 31bc6f07bea865039e470be6d4a8f0e2e9885ac8 Mon Sep 17 00:00:00 2001 From: naveensrinivasan <172697+naveensrinivasan@users.noreply.github.com> Date: Mon, 10 Jul 2023 15:23:05 -0500 Subject: [PATCH 2/2] Code review comments. Signed-off-by: naveensrinivasan <172697+naveensrinivasan@users.noreply.github.com> --- clients/githubrepo/roundtripper/rate_limit_test.go | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/clients/githubrepo/roundtripper/rate_limit_test.go b/clients/githubrepo/roundtripper/rate_limit_test.go index d0e17c976cb..b74b7d997b8 100644 --- a/clients/githubrepo/roundtripper/rate_limit_test.go +++ b/clients/githubrepo/roundtripper/rate_limit_test.go @@ -22,6 +22,7 @@ import ( ) func TestRoundTrip(t *testing.T) { + t.Parallel() var requestCount int ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { // Customize the response headers and body based on the test scenario @@ -38,7 +39,7 @@ func TestRoundTrip(t *testing.T) { w.Write([]byte("Success")) // nolint: errcheck } else { // First request: Return Retry-After header - w.Header().Set("Retry-After", "5") + w.Header().Set("Retry-After", "1") w.WriteHeader(http.StatusTooManyRequests) w.Write([]byte("Rate Limit Exceeded")) // nolint: errcheck } @@ -48,7 +49,9 @@ func TestRoundTrip(t *testing.T) { w.Write([]byte("Success")) // nolint: errcheck } })) - defer ts.Close() + t.Cleanup(func() { + defer ts.Close() + }) // Create the rateLimitTransport with the test server as the inner transport and a default logger transport := &rateLimitTransport{