Add tests for the rate limit retry logic.

Signed-off-by: David Calavera <[email protected]>

Author: calavera

go/Gopkg.lock

@@ -52,3 +52,7 @@
 [[constraint]]
   name = "github.com/Azure/go-autorest"
   version = "v9.6.0"
+
+[[constraint]]
+  name = "github.com/stretchr/testify"
+  version = "v1.1.4"

go/Gopkg.toml

@@ -26,45 +26,68 @@ func NewRetryableTransport(tr runtime.ClientTransport, attempts int) *RetryableT
 	}
 }
 
-func (tr *RetryableTransport) Submit(op *runtime.ClientOperation) (interface{}, error) {
-	op.Client.Transport = &retryableRoundTripper{
-		tr:       op.Client.Transport,
-		attempts: tr.attempts,
+func (t *RetryableTransport) Submit(op *runtime.ClientOperation) (interface{}, error) {
+	client := op.Client
+
+	if client == nil {
+		client = http.DefaultClient
+	}
+
+	transport := client.Transport
+	if transport == nil {
+		transport = http.DefaultTransport
 	}
+	client.Transport = &retryableRoundTripper{
+		tr:       transport,
+		attempts: t.attempts,
+	}
+
+	op.Client = client
 
-	return tr.Submit(op)
+	return t.tr.Submit(op)
 }
 
 func (tr *retryableRoundTripper) RoundTrip(req *http.Request) (resp *http.Response, err error) {
 	rr := autorest.NewRetriableRequest(req)
 
 	// Increment to add the first call (attempts denotes number of retries)
-	tr.attempts++
-	for attempt := 0; attempt < tr.attempts; attempt++ {
+	attempts := tr.attempts
+	attempts++
+	for attempt := 0; attempt < attempts; attempt++ {
 		err = rr.Prepare()
 		if err != nil {
 			return resp, err
 		}
-		resp, err = tr.RoundTrip(rr.Request())
+
+		resp, err = tr.tr.RoundTrip(rr.Request())
+
 		if err != nil || resp.StatusCode != http.StatusTooManyRequests {
 			return resp, err
 		}
-		delayWithRateLimit(resp, req.Cancel)
+
+		if !delayWithRateLimit(resp, req.Cancel) {
+			return resp, err
+		}
 	}
+
 	return resp, err
 }
 
-func delayWithRateLimit(resp *http.Response, cancel <-chan struct{}) {
-	retryReset, err := strconv.ParseInt(resp.Header.Get("X-RateLimit-Reset"), 10, 0)
+func delayWithRateLimit(resp *http.Response, cancel <-chan struct{}) bool {
+	r := resp.Header.Get("X-RateLimit-Reset")
+	if r == "" {
+		return false
+	}
+	retryReset, err := strconv.ParseInt(r, 10, 0)
 	if err != nil {
-		return
+		return false
 	}
 
 	t := time.Unix(retryReset, 0)
 	select {
 	case <-time.After(t.Sub(time.Now())):
-		return
+		return true
 	case <-cancel:
-		return
+		return false
 	}
 }

go/porcelain/http/http.go

@@ -0,0 +1,66 @@
+package http
+
+import (
+	"errors"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+	"time"
+
+	"github.com/go-openapi/runtime"
+	httptransport "github.com/go-openapi/runtime/client"
+	"github.com/go-openapi/strfmt"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRetryableTransport(t *testing.T) {
+	attempts := 0
+
+	server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		attempts++
+
+		if attempts == 1 {
+			reset := fmt.Sprintf("%d", time.Now().Add(1*time.Second).Unix())
+			rw.Header().Set("X-RateLimit-Reset", reset)
+			rw.WriteHeader(http.StatusTooManyRequests)
+			_, _ = rw.Write([]byte("rate limited"))
+		} else {
+			rw.WriteHeader(http.StatusOK)
+			_, _ = rw.Write([]byte("ok"))
+		}
+	}))
+	defer server.Close()
+
+	rwrtr := runtime.ClientRequestWriterFunc(func(req runtime.ClientRequest, _ strfmt.Registry) error {
+		return nil
+	})
+
+	hu, _ := url.Parse(server.URL)
+	rt := NewRetryableTransport(httptransport.New(hu.Host, "/", []string{"http"}), 2)
+
+	res, err := rt.Submit(&runtime.ClientOperation{
+		ID:          "getSites",
+		Method:      "GET",
+		PathPattern: "/",
+		Params:      rwrtr,
+		Reader: runtime.ClientResponseReaderFunc(func(response runtime.ClientResponse, consumer runtime.Consumer) (interface{}, error) {
+			if response.Code() == 200 {
+				var result string
+				if err := consumer.Consume(response.Body(), &result); err != nil {
+					return nil, err
+				}
+				return result, nil
+			}
+			return nil, errors.New("Generic error")
+		}),
+	})
+
+	require.NoError(t, err)
+	require.Equal(t, 2, attempts)
+
+	require.IsType(t, "", res)
+	actual := res.(string)
+	require.EqualValues(t, "ok", actual)
+}