refactor: split out middleware for later re-use

Also expose the burst parameter offered by `rate.NewLimiter`.

Author: jlucktay

src/mock/server_options.go

@@ -134,12 +134,15 @@ func WithRequestMatchPagesEnterprise(
 	return WithRequestMatchPages(ep, pages...)
 }
 
-// WithRateLimit enforces a rate limit using [golang.org/x/time/rate].
+// rateLimitMiddleware enforces a rate limit using [golang.org/x/time/rate].
 // rps is the number of requests per second allowed by the rate limiter.
-func WithRateLimit(rps float64) MockBackendOption {
-	limiter := rate.NewLimiter(rate.Limit(rps), 1)
+// Higher burst values allow more calls to happen at once.
+// A zero value for burst will not allow any events, unless rps == [rate.Inf].
+// This middleware is intended to be used with [github.com/gorilla/mux].
+func rateLimitMiddleware(rps float64, burst int) func(next http.Handler) http.Handler {
+	limiter := rate.NewLimiter(rate.Limit(rps), burst)
 
-	rateLimitMiddleware := func(next http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			if !limiter.Allow() {
 				// These values are based on this bit of logic within [github.com/google/go-github]:
@@ -153,8 +156,11 @@ func WithRateLimit(rps float64) MockBackendOption {
 			next.ServeHTTP(w, r)
 		})
 	}
+}
 
+// WithRateLimit enforces a rate limit on the mocked [http.Client].
+func WithRateLimit(rps float64, burst int) MockBackendOption {
 	return func(router *mux.Router) {
-		router.Use(rateLimitMiddleware)
+		router.Use(rateLimitMiddleware(rps, burst))
 	}
 }

src/mock/server_options_external_test.go

@@ -25,8 +25,10 @@ func TestWithRateLimit(t *testing.T) {
 			[]github.Repository{{Name: github.String(repoTwo)}},
 		),
 
-		// The rate limiter will allow 10 requests per second.
-		mock.WithRateLimit(10),
+		// The rate limiter will allow 10 requests per second, and a burst size of 1.
+		// These two options together mean that the rate of requests will be strictly enforced, so if any two requests are
+		// made less than 1/10th of a second apart, the latter will be refused and come back with a rate limit error.
+		mock.WithRateLimit(10, 1),
 	)
 
 	ghc := github.NewClient(mhc)
@@ -45,7 +47,8 @@ func TestWithRateLimit(t *testing.T) {
 
 			rleCount++
 
-			// Sleeping for one tenth of a second should be enough for the rate limiter on the mock server.
+			// After hitting the rate limiter and getting the appropriate error above, sleeping for one tenth of a second
+			// should be enough to reset the limiter and try requesting the repo list again.
 			time.Sleep(time.Second / 10)
 
 			continue