diff --git a/pool.go b/pool.go index 18c924a..0ec096f 100644 --- a/pool.go +++ b/pool.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "log/slog" + "net/http" "git.wisehodl.dev/jay/go-honeybee/transport" "git.wisehodl.dev/jay/go-honeybee/types" @@ -82,16 +83,20 @@ type PoolConfig struct { InboxBufferSize int EventsBufferSize int WorkerConfig WorkerConfig + RequestHeader http.Header } type PoolOption func(*PoolConfig) func NewPoolConfig(opts ...PoolOption) (*PoolConfig, error) { workerCfg, _ := NewWorkerConfig() + defaultHeader := http.Header{} + defaultHeader.Set("User-Agent", "honeybee/0.1.0") cfg := &PoolConfig{ InboxBufferSize: 256, EventsBufferSize: 10, WorkerConfig: *workerCfg, + RequestHeader: defaultHeader, } for _, o := range opts { o(cfg) @@ -120,6 +125,16 @@ func WithWorkerConfig(wc WorkerConfig) PoolOption { } } +func WithRequestHeader(h http.Header) PoolOption { + return func(c *PoolConfig) { + if h == nil { + c.RequestHeader = nil + } else { + c.RequestHeader = h.Clone() + } + } +} + func ValidatePoolConfig(c *PoolConfig) error { if c.InboxBufferSize < 1 { return fmt.Errorf("invalid inbox buffer size: %d", c.InboxBufferSize) @@ -328,9 +343,12 @@ func (p *Pool) Connect(id string, opts ...ConnectOption) error { opt(o) } - // TODO: pass headers + var hdr http.Header + if p.config.RequestHeader != nil { + hdr = p.config.RequestHeader.Clone() + } dialFn := func(ctx context.Context) (types.Socket, error) { - socket, _, err := p.dialer.DialContext(ctx, id, nil) + socket, _, err := p.dialer.DialContext(ctx, id, hdr) return socket, err } diff --git a/pool_test.go b/pool_test.go index 8bbe387..5aba45a 100644 --- a/pool_test.go +++ b/pool_test.go @@ -302,6 +302,85 @@ func TestPoolRetire(t *testing.T) { }) } +func TestPoolRequestHeader(t *testing.T) { + t.Run("default pool passes User-Agent header to DialContext", func(t *testing.T) { + config, _ := NewPoolConfig() + pool, _ := NewPool(context.Background(), config, nil) + + var capturedHeader http.Header + pool.dialer = &honeybeetest.MockDialer{ + DialContextFunc: func( + _ context.Context, _ string, h http.Header, + ) (types.Socket, *http.Response, error) { + capturedHeader = h + return honeybeetest.NewMockSocket(), nil, nil + }, + } + + pool.Connect("wss://test") + expectEvent(t, pool.events, "wss://test", EventConnected) + + assert.Equal(t, "honeybee/0.1.0", capturedHeader.Get("User-Agent")) + pool.Close() + }) + + t.Run("WithRequestHeader passes header to DialContext", func(t *testing.T) { + h := http.Header{} + h.Set("X-Custom", "value") + config, _ := NewPoolConfig(WithRequestHeader(h)) + pool, _ := NewPool(context.Background(), config, nil) + + var capturedHeader http.Header + pool.dialer = &honeybeetest.MockDialer{ + DialContextFunc: func( + _ context.Context, _ string, h http.Header, + ) (types.Socket, *http.Response, error) { + capturedHeader = h + return honeybeetest.NewMockSocket(), nil, nil + }, + } + + pool.Connect("wss://test") + expectEvent(t, pool.events, "wss://test", EventConnected) + + assert.Equal(t, "value", capturedHeader.Get("X-Custom")) + pool.Close() + }) + + t.Run("pool stores a clone - external mutation does not affect dial calls", func(t *testing.T) { + h := http.Header{} + h.Set("X-Custom", "original") + config, _ := NewPoolConfig(WithRequestHeader(h)) + + // mutate original after config is built + h.Set("X-Custom", "mutated") + + pool, _ := NewPool(context.Background(), config, nil) + + var capturedHeader http.Header + pool.dialer = &honeybeetest.MockDialer{ + DialContextFunc: func( + _ context.Context, _ string, h http.Header, + ) (types.Socket, *http.Response, error) { + capturedHeader = h + return honeybeetest.NewMockSocket(), nil, nil + }, + } + + pool.Connect("wss://test") + expectEvent(t, pool.events, "wss://test", EventConnected) + + assert.Equal(t, "original", capturedHeader.Get("X-Custom")) + pool.Close() + }) + + t.Run("WithRequestHeader(nil) stores nil cleanly", func(t *testing.T) { + cfg, err := NewPoolConfig(WithRequestHeader(nil)) + assert.NoError(t, err) + assert.Nil(t, cfg.RequestHeader) + }) +} + func TestPoolSend(t *testing.T) { mockSocket := honeybeetest.NewMockSocket() outgoingData := make(chan honeybeetest.MockOutgoingData, 10)