Adds tests for ws custom headers

pkg/ws/ws_test.go

@@ -2,6 +2,7 @@ package ws
 
 import (
 	"net/http/httptest"
+	"strings"
 	"testing"
 	"time"
 
@@ -80,4 +81,53 @@ func TestNewWSDisconnect(t *testing.T) {
 	case <-time.After(time.Millisecond * 10):
 		t.Errorf("Expected channel to be closed")
 	}
+
+	// Close again to make sure it doesn't panic
+	ws.Close()
+}
+
+func TestNewWSWithHeaders(t *testing.T) {
+	server := httptest.NewServer(websocket.Handler(func(ws *websocket.Conn) {
+		headerValue := ws.Request().Header.Get("X-Test")
+
+		if headerValue != "Test" {
+			t.Errorf("Expected header value to be 'Test', but got %v", headerValue)
+		}
+
+		var msg string
+		_ = websocket.Message.Receive(ws, &msg) // wait for request
+		_, _ = ws.Write([]byte(msg))
+		time.Sleep(time.Second) // to keep the connection open
+	}))
+	defer server.Close()
+
+	url := "ws://" + server.Listener.Addr().String()
+	ws, err := NewWS(url, Options{Headers: []string{"X-Test: Test"}})
+
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	ws.Close()
+}
+
+func TestNewWSWithInvalidHeaders(t *testing.T) {
+	server := httptest.NewServer(websocket.Handler(func(ws *websocket.Conn) {
+		var msg string
+		_ = websocket.Message.Receive(ws, &msg) // wait for request
+		_, _ = ws.Write([]byte(msg))
+		time.Sleep(time.Second) // to keep the connection open
+	}))
+	defer server.Close()
+
+	url := "ws://" + server.Listener.Addr().String()
+	_, err := NewWS(url, Options{Headers: []string{"X-Test"}})
+
+	if err == nil {
+		t.Fatalf("Expected error, but got nil")
+	}
+
+	if !strings.Contains(err.Error(), "invalid header") {
+		t.Errorf("Expected error to contain 'invalid header', but got %v", err)
+	}
 }