Add CheckRedirect option to http client (#498)

client/http/http_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"compress/flate"
 	"compress/gzip"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"net/http"
@@ -82,6 +83,25 @@ func TestTracedClient_Do(t *testing.T) {
 	}
 }
 
+func TestTracedClient_Do_Redirect(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		http.Redirect(w, r, "https://google.com", http.StatusSeeOther)
+	}))
+	defer ts.Close()
+	c, err := New(CheckRedirect(func(req *http.Request, via []*http.Request) error {
+		return errors.New("stop redirects")
+	}))
+	assert.NoError(t, err)
+	req, err := http.NewRequest("GET", ts.URL, nil)
+	assert.NoError(t, err)
+
+	res, err := c.Do(req)
+
+	assert.Errorf(t, err, "stop redirects")
+	assert.NotNil(t, res)
+	assert.Equal(t, http.StatusSeeOther, res.StatusCode)
+}
+
 func TestNew(t *testing.T) {
 	type args struct {
 		oo []OptionFunc
@@ -91,10 +111,16 @@ func TestNew(t *testing.T) {
 		args    args
 		wantErr bool
 	}{
-		{name: "success", args: args{oo: []OptionFunc{Timeout(time.Second), CircuitBreaker("test", circuitbreaker.Setting{}), Transport(&http.Transport{})}}, wantErr: false},
+		{name: "success", args: args{oo: []OptionFunc{
+			Timeout(time.Second),
+			CircuitBreaker("test", circuitbreaker.Setting{}),
+			Transport(&http.Transport{}),
+			CheckRedirect(func(req *http.Request, via []*http.Request) error { return nil }),
+		}}, wantErr: false},
 		{name: "failure, invalid timeout", args: args{oo: []OptionFunc{Timeout(0 * time.Second)}}, wantErr: true},
 		{name: "failure, invalid circuit breaker", args: args{[]OptionFunc{CircuitBreaker("", circuitbreaker.Setting{})}}, wantErr: true},
 		{name: "failure, invalid transport", args: args{[]OptionFunc{Transport(nil)}}, wantErr: true},
+		{name: "failure, invalid check redirect", args: args{[]OptionFunc{CheckRedirect(nil)}}, wantErr: true},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

client/http/option.go

@@ -46,3 +46,14 @@ func Transport(rt http.RoundTripper) OptionFunc {
 		return nil
 	}
 }
+
+// CheckRedirect option for setting the CheckRedirect for the client.
+func CheckRedirect(cr func(req *http.Request, via []*http.Request) error) OptionFunc {
+	return func(tc *TracedClient) error {
+		if cr == nil {
+			return errors.New("check redirect must be supplied")
+		}
+		tc.cl.CheckRedirect = cr
+		return nil
+	}
+}

client/http/option_test.go

@@ -2,6 +2,8 @@ package http
 
 import (
 	"net/http"
+	"reflect"
+	"runtime"
 	"testing"
 
 	"github.com/opentracing-contrib/go-stdlib/nethttp"
@@ -23,3 +25,24 @@ func TestTransport_Nil(t *testing.T) {
 	assert.Nil(t, client)
 	assert.Error(t, err, "transport must be supplied")
 }
+
+func TestCheckRedirect_Nil(t *testing.T) {
+	client, err := New(CheckRedirect(nil))
+
+	assert.Nil(t, client)
+	assert.Error(t, err, "check redirect must be supplied")
+}
+
+func TestCheckRedirect(t *testing.T) {
+	cr := func(req *http.Request, via []*http.Request) error {
+		return nil
+	}
+
+	client, err := New(CheckRedirect(cr))
+	assert.NoError(t, err)
+	assert.NotNil(t, client)
+
+	expFuncName := runtime.FuncForPC(reflect.ValueOf(cr).Pointer()).Name()
+	actFuncName := runtime.FuncForPC(reflect.ValueOf(client.cl.CheckRedirect).Pointer()).Name()
+	assert.Equal(t, expFuncName, actFuncName)
+}