Skip to content

Commit ef04be2

Browse files
🌱 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 <[email protected]>
1 parent 15b9046 commit ef04be2

2 files changed

Lines changed: 92 additions & 1 deletion

File tree

clients/githubrepo/roundtripper/rate_limit.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ type rateLimitTransport struct {
4242
innerTransport http.RoundTripper
4343
}
4444

45-
// Roundtrip handles caching and ratelimiting of responses from GitHub.
45+
// RoundTrip handles caching and rate-limiting of responses from GitHub.
4646
func (gh *rateLimitTransport) RoundTrip(r *http.Request) (*http.Response, error) {
4747
resp, err := gh.innerTransport.RoundTrip(r)
4848
if err != nil {
Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
// Copyright 2023 OpenSSF Scorecard Authors
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
package roundtripper
15+
16+
import (
17+
"net/http"
18+
"net/http/httptest"
19+
"testing"
20+
21+
"github.com/ossf/scorecard/v4/log"
22+
)
23+
24+
func TestRoundTrip(t *testing.T) {
25+
var requestCount int
26+
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
27+
// Customize the response headers and body based on the test scenario
28+
switch r.URL.Path {
29+
case "/error":
30+
w.WriteHeader(http.StatusInternalServerError)
31+
w.Write([]byte("Internal Server Error")) // nolint: errcheck
32+
case "/retry":
33+
requestCount++
34+
if requestCount == 2 {
35+
// Second request: Return successful response
36+
w.Header().Set("X-RateLimit-Remaining", "10")
37+
w.WriteHeader(http.StatusOK)
38+
w.Write([]byte("Success")) // nolint: errcheck
39+
} else {
40+
// First request: Return Retry-After header
41+
w.Header().Set("Retry-After", "5")
42+
w.WriteHeader(http.StatusTooManyRequests)
43+
w.Write([]byte("Rate Limit Exceeded")) // nolint: errcheck
44+
}
45+
case "/success":
46+
w.Header().Set("X-RateLimit-Remaining", "10")
47+
w.WriteHeader(http.StatusOK)
48+
w.Write([]byte("Success")) // nolint: errcheck
49+
}
50+
}))
51+
defer ts.Close()
52+
53+
// Create the rateLimitTransport with the test server as the inner transport and a default logger
54+
transport := &rateLimitTransport{
55+
innerTransport: ts.Client().Transport,
56+
logger: log.NewLogger(log.DefaultLevel),
57+
}
58+
59+
t.Run("Successful response", func(t *testing.T) {
60+
req, err := http.NewRequest(http.MethodGet, ts.URL+"/success", nil)
61+
if err != nil {
62+
t.Fatalf("Failed to create request: %v", err)
63+
}
64+
65+
resp, err := transport.RoundTrip(req)
66+
if err != nil {
67+
t.Errorf("Unexpected error: %v", err)
68+
}
69+
if resp.StatusCode != http.StatusOK {
70+
t.Errorf("Expected status code %d, got %d", http.StatusOK, resp.StatusCode)
71+
}
72+
})
73+
74+
t.Run("Retry-After header set", func(t *testing.T) {
75+
req, err := http.NewRequest(http.MethodGet, ts.URL+"/retry", nil)
76+
if err != nil {
77+
t.Fatalf("Failed to create request: %v", err)
78+
}
79+
80+
resp, err := transport.RoundTrip(req)
81+
if err != nil {
82+
t.Errorf("Unexpected error: %v", err)
83+
}
84+
if resp.StatusCode != http.StatusOK {
85+
t.Errorf("Expected status code %d, got %d", http.StatusOK, resp.StatusCode)
86+
}
87+
if requestCount != 2 {
88+
t.Errorf("Expected 2 requests, got %d", requestCount)
89+
}
90+
})
91+
}

0 commit comments

Comments
 (0)