diff --git a/pkg/httpsign/options.go b/pkg/httpsign/options.go new file mode 100644 index 0000000..b8aae03 --- /dev/null +++ b/pkg/httpsign/options.go @@ -0,0 +1,17 @@ +package httpsign + +import ( + "net/http" + "time" +) + +type Option func(c *http.Client) + +func WithTimeout(timeout time.Duration) Option { + return func(c *http.Client) { + if timeout == 0 { + return + } + c.Timeout = timeout + } +} diff --git a/pkg/httpsign/signclient.go b/pkg/httpsign/signclient.go index f5fc994..6e3b613 100644 --- a/pkg/httpsign/signclient.go +++ b/pkg/httpsign/signclient.go @@ -4,10 +4,20 @@ import ( "net/http" ) -func NewClient(transport http.RoundTripper, key string, secret []byte) *http.Client { - return &http.Client{ +func NewClient(transport http.RoundTripper, key string, secret []byte, opts ...Option) *http.Client { + c := &http.Client{ Transport: NewTransport(transport, key, secret), } + + for i := range opts { + if opts[i] == nil { + continue + } + + opts[i](c) + } + + return c } func NewTransport(inner http.RoundTripper, key string, secret []byte) HTTPSign { diff --git a/pkg/httpsign/signclient_test.go b/pkg/httpsign/signclient_test.go index 1010a02..2fcc490 100644 --- a/pkg/httpsign/signclient_test.go +++ b/pkg/httpsign/signclient_test.go @@ -4,8 +4,10 @@ import ( "io" "net/http" "testing" + "time" "github.com/KyberNetwork/tradinglib/pkg/httpsign" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -24,3 +26,12 @@ func TestRequest(t *testing.T) { require.NoError(t, err) t.Log(string(out)) } + +func TestOption(t *testing.T) { + timeout := time.Second * 3 + + client := httpsign.NewClient( + http.DefaultTransport, "..", []byte(".."), httpsign.WithTimeout(timeout), + ) + assert.Equal(t, client.Timeout, timeout) +}