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_OnConnectErrorTriggersRedial 断言连接钩子失败会触发重拨：等到第二次
// 拨号事件发生后才取消上下文，而不是先睡固定的 60ms 再看拨了几次——后者在并行
// 负载下可能只够拨一次，等于把断言交给机器负载决定。
func TestClient_OnConnectErrorTriggersRedial(t *testing.T) {
	const (
		// redialTarget 是断言「发生过重拨」所需的拨号次数：1 次首拨 + 1 次重拨。
		redialTarget = 2
		// dialEvents 是拨号事件通道的容量，保证 Run 在测试取消之前不会阻塞在投递上。
		dialEvents = 8
	)

	var dials atomic.Int64
	// 每次拨号投递一个事件；测试等满 redialTarget 次再取消，因此断言的是
	// 「重拨确实发生了」这件事本身，与墙钟无关。
	dialed := make(chan struct{}, dialEvents)
	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 恒错，因此每次会话都被客户端拒绝、拨号反复发生。
	client := New(
		func(context.Context) (string, error) {
			dials.Add(1)
			dialed <- struct{}{}
			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())
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- client.Run(ctx) }()

	for i := 0; i < redialTarget; i++ {
		select {
		case <-dialed:
		case <-time.After(3 * time.Second):
			t.Fatalf("timed out waiting for dial %d", i+1)
		}
	}
	cancel()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("Run did not exit")
	}
	if dials.Load() < redialTarget {
		t.Errorf("dials = %d, want >= %d (reconnect on failed connect hook)", dials.Load(), redialTarget)
	}
}