Move websocket headers to opt function 'WithWebsocketHeaders'

graphql/client.go

@@ -133,16 +133,10 @@ type WebSocketOption func(*webSocketClient)
 //
 // The client does not support queries nor mutations, and will return an error
 // if passed a request that attempts one.
-func NewClientUsingWebSocket(endpoint string, wsDialer Dialer, headers http.Header, opts ...WebSocketOption) WebSocketClient {
-	if headers == nil {
-		headers = http.Header{}
-	}
-	if headers.Get("Sec-WebSocket-Protocol") == "" {
-		headers.Add("Sec-WebSocket-Protocol", "graphql-transport-ws")
-	}
+func NewClientUsingWebSocket(endpoint string, wsDialer Dialer, opts ...WebSocketOption) WebSocketClient {
 	client := &webSocketClient{
 		Dialer:        wsDialer,
-		Header:        headers,
+		Header:        http.Header{},
 		errChan:       make(chan error),
 		endpoint:      endpoint,
 		subscriptions: subscriptionMap{map_: make(map[string]subscription)},
@@ -152,6 +146,10 @@ func NewClientUsingWebSocket(endpoint string, wsDialer Dialer, headers http.Head
 		opt(client)
 	}
 
+	if client.Header.Get("Sec-WebSocket-Protocol") == "" {
+		client.Header.Add("Sec-WebSocket-Protocol", "graphql-transport-ws")
+	}
+
 	return client
 }
 
@@ -163,6 +161,13 @@ func WithConnectionParams(connParams map[string]interface{}) WebSocketOption {
 	}
 }
 
+// WithWebsocketHeaders sets the headers to be sent to the server.
+func WithWebsocketHeaders(headers http.Header) WebSocketOption {
+	return func(ws *webSocketClient) {
+		ws.Header = headers
+	}
+}
+
 func newClient(endpoint string, httpClient Doer, method string) Client {
 	if httpClient == nil || httpClient == (*http.Client)(nil) {
 		httpClient = http.DefaultClient

internal/integration/roundtrip.go

@@ -167,7 +167,6 @@ func newRoundtripWebSocketClient(t *testing.T, endpoint string, opts ...graphql.
 		wsWrapped: graphql.NewClientUsingWebSocket(
 			endpoint,
 			&MyDialer{Dialer: dialer},
-			nil,
 			opts...,
 		),
 		t: t,