Merge 81220a7 into 265a40c

pkg/drivers/http/driver.go

@@ -3,13 +3,14 @@ package http
 import (
 	"bytes"
 	"context"
-	"github.com/MontFerret/ferret/pkg/runtime/logging"
-	"github.com/MontFerret/ferret/pkg/runtime/values"
-	"github.com/gobwas/glob"
 	"io"
 	"net/http"
 	"net/url"
 
+	"github.com/MontFerret/ferret/pkg/runtime/logging"
+	"github.com/MontFerret/ferret/pkg/runtime/values"
+	"github.com/gobwas/glob"
+
 	"golang.org/x/net/html/charset"
 
 	"github.com/PuerkitoBio/goquery"
@@ -32,16 +33,18 @@ func NewDriver(opts ...Option) *Driver {
 	drv.options = NewOptions(opts)
 
 	drv.client = newHTTPClient(drv.options)
-	drv.client.Concurrency = drv.options.Concurrency
-	drv.client.MaxRetries = drv.options.MaxRetries
-	drv.client.Backoff = drv.options.Backoff
 
 	return drv
 }
 
 func newHTTPClient(options *Options) (httpClient *pester.Client) {
 	httpClient = pester.New()
 
+	httpClient.Concurrency = options.Concurrency
+	httpClient.MaxRetries = options.MaxRetries
+	httpClient.Backoff = options.Backoff
+	httpClient.Timeout = options.Timeout
+
 	if options.HTTPTransport != nil {
 		httpClient.Transport = options.HTTPTransport
 	}
@@ -54,8 +57,6 @@ func newHTTPClient(options *Options) (httpClient *pester.Client) {
 		return
 	}
 
-	httpClient = pester.NewExtendedClient(&http.Client{Transport: httpClient.Transport})
-
 	return
 }
 
@@ -219,7 +220,7 @@ func (drv *Driver) makeRequest(ctx context.Context, req *http.Request, params dr
 		params.Headers.ForEach(func(value []string, key string) bool {
 			v := params.Headers.Get(key)
 
-			req.Header.Add(key, v)
+			req.Header.Set(key, v)
 
 			logger.
 				Debug().

pkg/drivers/http/driver_test.go

@@ -39,15 +39,6 @@ func Test_newHTTPClientWithTransport(t *testing.T) {
 				HTTPTransport: httpTransport,
 			}},
 		},
-		{
-			name: "check transport exist with pester.NewExtendedClient()",
-			args: args{options: &Options{
-				Options: &drivers.Options{
-					Proxy: "http://0.0.0.0",
-				},
-				HTTPTransport: httpTransport,
-			}},
-		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -95,24 +86,6 @@ func Test_newHTTPClient(t *testing.T) {
 
 		So(hc, ShouldBeNil)
 	})
-
-	Convey("pester.NewExtend()", t, func() {
-		var (
-			client = newHTTPClient(&Options{
-				Options: &drivers.Options{
-					Proxy: "http://0.0.0.0",
-				},
-			})
-
-			rValue = reflect.ValueOf(client).Elem()
-			rField = rValue.Field(0)
-		)
-
-		rField = reflect.NewAt(rField.Type(), unsafe.Pointer(rField.UnsafeAddr())).Elem()
-		hc := rField.Interface().(*http.Client)
-
-		So(hc, ShouldNotBeNil)
-	})
 }
 
 func TestDriver_convertToUTF8(t *testing.T) {

pkg/drivers/http/options.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	stdhttp "net/http"
+	"time"
 
 	"github.com/gobwas/glob"
 	"github.com/sethgrid/pester"
@@ -12,6 +13,7 @@ import (
 var (
 	DefaultConcurrency = 1
 	DefaultMaxRetries  = 5
+	DefaultTimeout     = time.Second * 30
 )
 
 type (
@@ -29,6 +31,7 @@ type (
 		Concurrency     int
 		HTTPCodesFilter []compiledStatusCodeFilter
 		HTTPTransport   *stdhttp.Transport
+		Timeout         time.Duration
 	}
 )
 
@@ -39,6 +42,7 @@ func NewOptions(setters []Option) *Options {
 	opts.Backoff = pester.ExponentialBackoff
 	opts.Concurrency = DefaultConcurrency
 	opts.MaxRetries = DefaultMaxRetries
+	opts.Timeout = DefaultTimeout
 	opts.HTTPCodesFilter = make([]compiledStatusCodeFilter, 0, 5)
 
 	for _, setter := range setters {
@@ -143,3 +147,9 @@ func WithCustomTransport(transport *stdhttp.Transport) Option {
 		opts.HTTPTransport = transport
 	}
 }
+
+func WithTimeout(duration time.Duration) Option {
+	return func(opts *Options) {
+		opts.Timeout = duration
+	}
+}

pkg/drivers/http/options_test.go

@@ -30,6 +30,7 @@ func TestNewOptions(t *testing.T) {
 		expectedMaxRetries := 2
 		expectedConcurrency := 10
 		expectedTransport := &stdhttp.Transport{}
+		expectedTimeout := time.Second * 5
 
 		opts := http.NewOptions([]http.Option{
 			http.WithCustomName(expectedName),
@@ -69,6 +70,7 @@ func TestNewOptions(t *testing.T) {
 			http.WithAllowedHTTPCode(401),
 			http.WithAllowedHTTPCodes([]int{403, 404}),
 			http.WithCustomTransport(expectedTransport),
+			http.WithTimeout(time.Second * 5),
 		})
 		So(opts.Options, ShouldNotBeNil)
 		So(opts.Name, ShouldEqual, expectedName)
@@ -81,5 +83,6 @@ func TestNewOptions(t *testing.T) {
 		So(opts.Concurrency, ShouldEqual, expectedConcurrency)
 		So(opts.HTTPCodesFilter, ShouldHaveLength, 3)
 		So(opts.HTTPTransport, ShouldEqual, expectedTransport)
+		So(opts.Timeout, ShouldEqual, expectedTimeout)
 	})
 }