im_server/loadtest/im-pressure/wsclient_test.go

package main

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

func TestWSClientObservesCtrlAfterCommandTimeout(t *testing.T) {
	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()
		var command map[string]any
		if err = conn.ReadJSON(&command); err != nil {
			return
		}
		time.Sleep(40 * time.Millisecond)
		_ = conn.WriteJSON(map[string]any{"ctrl": map[string]any{"id": "pub-a-0001", "code": 202, "text": "accepted"}})
	}))
	defer server.Close()

	conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil)
	if err != nil {
		t.Fatal(err)
	}
	observed := make(chan ctrlObservation, 1)
	client := &WSClient{
		account: AccountRecord{Machine: "a", Index: 1}, conn: conn, commandTimeout: 10 * time.Millisecond,
		waits: make(map[string]ctrlWait), timedOutCtrls: make(map[string]timedOutCtrl), done: make(chan struct{}),
		onCtrlObserved: func(_ AccountRecord, observation ctrlObservation) { observed <- observation },
	}
	go client.readLoop()
	defer client.Close()

	_, err = client.command(context.Background(), "pub", map[string]any{"pub": map[string]any{"id": "pub-a-0001"}})
	if err == nil || !strings.Contains(err.Error(), "ctrl timeout") {
		t.Fatalf("command error=%v, want ctrl timeout", err)
	}
	select {
	case observation := <-observed:
		if !observation.Late || observation.Kind != "pub" || observation.Code != 202 || observation.Latency < 10*time.Millisecond {
			t.Fatalf("unexpected observation: %+v", observation)
		}
	case <-time.After(time.Second):
		t.Fatal("late ctrl response was not observed")
	}
}

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]ctrlWait), timedOutCtrls: make(map[string]timedOutCtrl), 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")
	}
}

func TestWSClientReadLoopClosesConnectionAfterDecodeError(t *testing.T) {
	serverSawClose := 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()
		if err = conn.WriteMessage(websocket.TextMessage, []byte("{")); err != nil {
			return
		}
		if _, _, err = conn.ReadMessage(); err != nil {
			serverSawClose <- struct{}{}
		}
	}))
	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]ctrlWait), timedOutCtrls: make(map[string]timedOutCtrl), done: make(chan struct{})}
	go client.readLoop()

	select {
	case <-client.done:
	case <-time.After(time.Second):
		t.Fatal("read loop did not exit after malformed frame")
	}
	select {
	case <-serverSawClose:
	case <-time.After(time.Second):
		t.Fatal("read loop exited without closing the underlying connection")
	}
}

func TestWSClientDataCallbackDoesNotBlockPong(t *testing.T) {
	pong := 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.SetPongHandler(func(string) error {
			select {
			case pong <- struct{}{}:
			default:
			}
			return nil
		})
		if err = conn.WriteJSON(map[string]any{"data": map[string]any{
			"topic": "grp-test", "seq": 1, "head": map[string]any{}, "content": map[string]any{},
		}}); err != nil {
			return
		}
		if err = conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(time.Second)); err != nil {
			return
		}
		conn.SetReadDeadline(time.Now().Add(time.Second))
		for {
			if _, _, readErr := conn.ReadMessage(); readErr != nil {
				return
			}
		}
	}))
	defer server.Close()

	conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil)
	if err != nil {
		t.Fatal(err)
	}
	blocked := make(chan struct{})
	client := &WSClient{
		conn: conn, waits: make(map[string]ctrlWait), timedOutCtrls: make(map[string]timedOutCtrl), done: make(chan struct{}),
		onData: func(AccountRecord, dataPacket) { <-blocked },
	}
	client.startDataProcessor()
	go client.readLoop()
	defer func() {
		close(blocked)
		client.Close()
	}()

	select {
	case <-pong:
	case <-time.After(2 * time.Second):
		t.Fatal("pong was blocked by the data callback")
	}
}