diff --git a/config.go b/config.go index 01f604f..0474a1b 100644 --- a/config.go +++ b/config.go @@ -2,6 +2,7 @@ package honeybee import ( "git.wisehodl.dev/jay/go-honeybee/transport" + "git.wisehodl.dev/jay/go-honeybee/types" "time" ) @@ -17,6 +18,7 @@ type PoolConfig struct { ConnectionConfig transport.ConnectionConfig WorkerFactory WorkerFactory WorkerConfig WorkerConfig + Dialer types.Dialer } type PoolOption func(*PoolConfig) error @@ -122,6 +124,13 @@ func WithWorkerConfig(wc WorkerConfig) PoolOption { } } +func WithPoolDialer(d types.Dialer) PoolOption { + return func(c *PoolConfig) error { + c.Dialer = d + return nil + } +} + func WithWorkerFactory(wf WorkerFactory) PoolOption { return func(c *PoolConfig) error { c.WorkerFactory = wf diff --git a/config_pool_test.go b/config_pool_test.go index 2010e79..d7cedc6 100644 --- a/config_pool_test.go +++ b/config_pool_test.go @@ -1,6 +1,7 @@ package honeybee import ( + "git.wisehodl.dev/jay/go-honeybee/honeybeetest" "git.wisehodl.dev/jay/go-honeybee/transport" "github.com/stretchr/testify/assert" "testing" @@ -138,3 +139,10 @@ func TestValidatePoolConfig(t *testing.T) { }) } } + +func TestWithPoolDialer(t *testing.T) { + mock := &honeybeetest.MockDialer{} + conf, err := NewPoolConfig(WithPoolDialer(mock)) + assert.NoError(t, err) + assert.Equal(t, mock, conf.Dialer) +} diff --git a/transport/config.go b/transport/config.go index f9599b0..e085038 100644 --- a/transport/config.go +++ b/transport/config.go @@ -1,6 +1,7 @@ package transport import ( + "git.wisehodl.dev/jay/go-honeybee/types" "net/http" "time" ) @@ -21,6 +22,7 @@ type ConnectionConfig struct { IncomingBufferSize int ErrorsBufferSize int Retry RetryConfig + Dialer types.Dialer } type RetryConfig struct { @@ -220,6 +222,13 @@ func WithErrorsBufferSize(value int) ConnectionOption { } } +func WithConnectionDialer(d types.Dialer) ConnectionOption { + return func(c *ConnectionConfig) error { + c.Dialer = d + return nil + } +} + func WithRetryDisabled() ConnectionOption { return func(c *ConnectionConfig) error { c.Retry.Disabled = true diff --git a/transport/config_test.go b/transport/config_test.go index b8acad2..7a0d921 100644 --- a/transport/config_test.go +++ b/transport/config_test.go @@ -1,6 +1,7 @@ package transport import ( + "git.wisehodl.dev/jay/go-honeybee/honeybeetest" "github.com/stretchr/testify/assert" "net/http" "testing" @@ -253,3 +254,10 @@ func TestValidateConnectionConfig(t *testing.T) { }) } } + +func TestWithConnectionDialer(t *testing.T) { + mock := &honeybeetest.MockDialer{} + conf, err := NewConnectionConfig(WithConnectionDialer(mock)) + assert.NoError(t, err) + assert.Equal(t, mock, conf.Dialer) +}