Cyrene_v2/internal/channel/adapter/wsutil/client_test.go

package wsutil

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/gorilla/websocket"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

var wsu = websocket.Upgrader{}

// TestClient_DialReadReconnect exercises the connect hook, the read loop, and the
// reconnect path: the server accepts a connection, sends one frame, drops it;
// the client re-reads a second frame on the next connection, then the test cancels
// the context and the client exits cleanly.
func TestClient_DialReadReconnect(t *testing.T) {
	var conns atomic.Int64
	var mu sync.Mutex
	var got []string
	received := make(chan string, 2)

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := wsu.Upgrade(w, r, nil)
		if err != nil {
			t.Errorf("upgrade: %v", err)
			return
		}
		defer conn.Close()
		n := conns.Add(1)
		msg := "msg-" + string(rune('0'+n))
		_ = conn.WriteMessage(websocket.TextMessage, []byte(msg))
		if n < 2 {
			return // close the first connection to trigger a reconnect
		}
		// Hold the second connection open until the client cancels the context.
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				return
			}
		}
	}))
	defer srv.Close()

	url := "ws://" + strings.TrimPrefix(srv.URL, "http://") + "/testws"
	client := New(
		func(context.Context) (string, error) { return url, nil },
		Handler{
			OnConnect: func(ctx context.Context, c *Conn) error { return nil },
			OnMessage: func(ctx context.Context, p []byte) error {
				mu.Lock()
				got = append(got, string(p))
				mu.Unlock()
				received <- string(p)
				return nil
			},
		},
		WithBackoff(NewBackoff(10*time.Millisecond, 50*time.Millisecond)),
		WithLogger(logger.New(logger.Options{Out: io.Discard})),
	)

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- client.Run(ctx) }()

	for i := 0; i < 2; i++ {
		select {
		case m := <-received:
			if m != "msg-" + string(rune('1'+i)) {
				t.Errorf("message[%d] = %q", i, m)
			}
		case <-time.After(3 * time.Second):
			t.Fatalf("timed out waiting for message %d", i+1)
		}
	}
	cancel()
	select {
	case err := <-done:
		if err == nil {
			t.Error("Run should return a non-nil error on cancel")
		}
	case <-time.After(3 * time.Second):
		t.Fatal("Run did not exit after cancel")
	}

	mu.Lock()
	defer mu.Unlock()
	if len(got) != 2 {
		t.Fatalf("received %d messages, want 2", len(got))
	}
	if conns.Load() < 2 {
		t.Errorf("connection count = %d, want >= 2 (reconnect)", conns.Load())
	}
}

// TestClient_HeartbeatVerifiesPing sends an s=2-like heartbeat every 20ms and
// asserts the dial URL is still the one provided.
func TestClient_OnConnectErrorTriggersRedial(t *testing.T) {
	var dials atomic.Int64
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := wsu.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer conn.Close()
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				return
			}
		}
	}))
	defer srv.Close()
	url := "ws://" + strings.TrimPrefix(srv.URL, "http://") + "/testws"

	// OnConnect always errors so the client keeps rejecting each session and the
	// dial fires repeatedly.
	client := New(
		func(context.Context) (string, error) { dials.Add(1); return url, nil },
		Handler{
			OnConnect: func(ctx context.Context, c *Conn) error { return context.DeadlineExceeded },
		},
		WithBackoff(NewBackoff(5*time.Millisecond, 20*time.Millisecond)),
		WithLogger(logger.New(logger.Options{Out: io.Discard})),
	)
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- client.Run(ctx) }()

	time.Sleep(60 * time.Millisecond)
	cancel()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("Run did not exit")
	}
	if dials.Load() < 2 {
		t.Errorf("dials = %d, want >= 2 (reconnect on failed connect hook)", dials.Load())
	}
}