Golang WebSocket 客户端测试库 wstest 使用指南

wstest 是一个用于测试 WebSocket 客户端的 Go 库，它可以帮助你模拟 WebSocket 服务器行为，方便进行单元测试。下面我将详细介绍如何使用 wstest 来测试 WebSocket 客户端。

安装 wstest

首先，使用 go get 安装 wstest：

go get github.com/gorilla/websocket
go get github.com/posener/wstest

基本使用方法

1. 最简单的测试用例

package main

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

	"github.com/gorilla/websocket"
	"github.com/posener/wstest"
	"github.com/stretchr/testify/assert"
)

func TestWebSocketClient(t *testing.T) {
	// 创建一个测试用的 WebSocket 服务器
	d := wstest.NewDialer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// 升级为 WebSocket 连接
		conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		if err != nil {
			t.Fatalf("升级连接失败: %v", err)
		}
		defer conn.Close()

		// 读取客户端消息
		_, msg, err := conn.ReadMessage()
		if err != nil {
			t.Fatalf("读取消息失败: %v", err)
		}

		// 回显消息
		if err := conn.WriteMessage(websocket.TextMessage, msg); err != nil {
			t.Fatalf("写入消息失败: %v", err)
		}
	}))

	// 客户端连接
	conn, resp, err := d.Dial("ws://example.org/ws", nil)
	if err != nil {
		t.Fatalf("连接失败: %v", err)
	}
	defer conn.Close()

	// 检查响应状态码
	assert.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode)

	// 发送测试消息
	testMsg := "hello, world"
	if err := conn.WriteMessage(websocket.TextMessage, []byte(testMsg)); err != nil {
		t.Fatalf("发送消息失败: %v", err)
	}

	// 读取回显消息
	_, msg, err := conn.ReadMessage()
	if err != nil {
		t.Fatalf("接收消息失败: %v", err)
	}

	// 验证消息
	assert.Equal(t, testMsg, string(msg))
}

2. 更复杂的测试场景

func TestWebSocketWithAuth(t *testing.T) {
	// 创建带认证的测试服务器
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// 检查认证头
		if r.Header.Get("Authorization") != "Bearer valid-token" {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}

		conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		if err != nil {
			t.Fatalf("升级连接失败: %v", err)
		}
		defer conn.Close()

		// 处理消息循环
		for {
			mt, msg, err := conn.ReadMessage()
			if err != nil {
				return
			}

			// 处理不同类型的消息
			switch mt {
			case websocket.TextMessage:
				conn.WriteMessage(websocket.TextMessage, []byte("echo: "+string(msg)))
			case websocket.BinaryMessage:
				conn.WriteMessage(websocket.BinaryMessage, msg)
			case websocket.PingMessage:
				conn.WriteMessage(websocket.PongMessage, nil)
			}
		}
	})

	d := wstest.NewDialer(handler)

	// 创建带认证头的请求
	headers := http.Header{}
	headers.Set("Authorization", "Bearer valid-token")

	// 连接服务器
	conn, _, err := d.Dial("ws://example.org/ws", headers)
	if err != nil {
		t.Fatalf("连接失败: %v", err)
	}
	defer conn.Close()

	// 测试文本消息
	textMsg := "test message"
	if err := conn.WriteMessage(websocket.TextMessage, []byte(textMsg)); err != nil {
		t.Fatalf("发送文本消息失败: %v", err)
	}

	_, resp, err := conn.ReadMessage()
	if err != nil {
		t.Fatalf("接收响应失败: %v", err)
	}
	assert.Equal(t, "echo: "+textMsg, string(resp))

	// 测试二进制消息
	binaryMsg := []byte{0x01, 0x02, 0x03}
	if err := conn.WriteMessage(websocket.BinaryMessage, binaryMsg); err != nil {
		t.Fatalf("发送二进制消息失败: %v", err)
	}

	_, respBinary, err := conn.ReadMessage()
	if err != nil {
		t.Fatalf("接收二进制响应失败: %v", err)
	}
	assert.Equal(t, binaryMsg, respBinary)
}

3. 测试连接关闭

func TestWebSocketClose(t *testing.T) {
	// 创建会主动关闭连接的服务器
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		if err != nil {
			t.Fatalf("升级连接失败: %v", err)
		}
		defer conn.Close()

		// 读取一条消息后关闭连接
		_, _, err = conn.ReadMessage()
		if err != nil {
			return
		}

		// 发送关闭帧
		conn.WriteMessage(websocket.CloseMessage, 
			websocket.FormatCloseMessage(websocket.CloseNormalClosure, "goodbye"))
	})

	d := wstest.NewDialer(handler)
	conn, _, err := d.Dial("ws://example.org/ws", nil)
	if err != nil {
		t.Fatalf("连接失败: %v", err)
	}
	defer conn.Close()

	// 发送消息触发服务器关闭
	if err := conn.WriteMessage(websocket.TextMessage, []byte("trigger close")); err != nil {
		t.Fatalf("发送消息失败: %v", err)
	}

	// 等待关闭
	_, _, err = conn.ReadMessage()
	assert.Error(t, err)
	assert.True(t, websocket.IsCloseError(err, websocket.CloseNormalClosure))
}

高级用法

1. 测试并发连接

func TestConcurrentWebSocketConnections(t *testing.T) {
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		if err != nil {
			t.Fatalf("升级连接失败: %v", err)
		}
		defer conn.Close()

		// 简单回显
		for {
			_, msg, err := conn.ReadMessage()
			if err != nil {
				return
			}
			conn.WriteMessage(websocket.TextMessage, msg)
		}
	})

	d := wstest.NewDialer(handler)

	// 测试并发连接
	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()

			conn, _, err := d.Dial("ws://example.org/ws", nil)
			if err != nil {
				t.Errorf("连接 %d 失败: %v", i, err)
				return
			}
			defer conn.Close()

			msg := fmt.Sprintf("message %d", i)
			if err := conn.WriteMessage(websocket.TextMessage, []byte(msg)); err != nil {
				t.Errorf("发送消息 %d 失败: %v", i, err)
				return
			}

			_, resp, err := conn.ReadMessage()
			if err != nil {
				t.Errorf("接收响应 %d 失败: %v", i, err)
				return
			}

			assert.Equal(t, msg, string(resp))
		}(i)
	}
	wg.Wait()
}

2. 测试超时场景

func TestWebSocketTimeout(t *testing.T) {
	// 创建慢速服务器
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		if err != nil {
			t.Fatalf("升级连接失败: %v", err)
		}
		defer conn.Close()

		// 慢速处理
		time.Sleep(2 * time.Second)
		conn.WriteMessage(websocket.TextMessage, []byte("response"))
	})

	d := wstest.NewDialer(handler)
	conn, _, err := d.Dial("ws://example.org/ws", nil)
	if err != nil {
		t.Fatalf("连接失败: %v", err)
	}
	defer conn.Close()

	// 设置读取超时
	conn.SetReadDeadline(time.Now().Add(1 * time.Second))

	// 发送请求
	if err := conn.WriteMessage(websocket.TextMessage, []byte("request")); err != nil {
		t.Fatalf("发送消息失败: %v", err)
	}

	// 读取响应应该会超时
	_, _, err = conn.ReadMessage()
	assert.Error(t, err)
	assert.True(t, os.IsTimeout(err))
}

总结

wstest 提供了一种简单的方式来测试 WebSocket 客户端代码，主要优点包括：

1. 不需要实际启动 HTTP 服务器
2. 可以完全控制服务器行为
3. 支持测试各种 WebSocket 协议特性
4. 易于集成到现有的测试框架中

通过合理使用 wstest，你可以为 WebSocket 客户端代码编写全面的单元测试，覆盖正常流程、错误处理和边界情况。