fix: keep load test websocket sessions alive

loadtest/im-pressure/wsclient_test.go

@@ -0,0 +1,50 @@
+package main
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+func TestWSClientKeepaliveSendsPing(t *testing.T) {
+	ping := make(chan struct{}, 1)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		defer conn.Close()
+		conn.SetPingHandler(func(string) error {
+			select {
+			case ping <- struct{}{}:
+			default:
+			}
+			return conn.WriteControl(websocket.PongMessage, nil, time.Now().Add(time.Second))
+		})
+		for {
+			if _, _, err = conn.ReadMessage(); err != nil {
+				return
+			}
+		}
+	}))
+	defer server.Close()
+
+	conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	client := &WSClient{conn: conn, waits: make(map[string]chan ctrlPacket), done: make(chan struct{})}
+	go client.readLoop()
+	client.startKeepalive(10 * time.Millisecond)
+	defer client.Close()
+
+	select {
+	case <-ping:
+	case <-time.After(time.Second):
+		t.Fatal("websocket keepalive ping was not received")
+	}
+}