如何为Gorilla Websocket的socketHandler编写单元测试用例

核心思路 #

单元测试不需要启动main函数里的对外监听服务，直接用Go标准库net/http/httptest包启动绑定socketHandler的内存级测试服务，再用Gorilla Websocket客户端直连这个测试服务，模拟各种请求场景验证逻辑即可，全程自动分配空闲端口，无环境依赖，执行速度快。
需要覆盖三类核心逻辑：

• 正常连接下的回声消息收发逻辑是否符合预期
• Websocket握手升级失败时服务是否正常返回错误，无panic
• 客户端主动断开连接时服务端是否正常退出循环，无异常

前置适配 #

原代码中使用默认配置的websocket.Upgrader，默认会校验请求的Origin头，非同源请求会直接拒绝升级。单元测试发起的请求不满足同源校验规则，需要在测试初始化阶段临时放开校验。如果不想测试过程中打印大量业务日志，也可以在初始化阶段把日志输出重定向丢弃。

具体测试用例实现 #

1. 正常消息回声测试 #

• 启动httptest测试服务，绑定socketHandler，测试结束自动关闭服务
• 将测试服务的HTTP地址替换为WS协议前缀，得到Websocket连接地址
• 用Websocket默认拨号器建立客户端连接，测试结束自动关闭客户端
• 构造多组测试用例（文本消息、二进制消息、空消息、长文本），依次发送给服务端，读取服务端返回的内容，校验消息类型、消息内容和发送的内容完全一致
• 所有读写操作设置2秒超时，避免异常场景下测试卡死

2. 握手升级失败测试 #

• 启动测试服务后，直接发起普通HTTP GET请求（不携带Websocket握手头）
• 校验响应状态码是否为400 Bad Request，确认服务端正确识别非法握手请求，无panic

3. 客户端断开测试 #

• 建立正常Websocket连接后，客户端主动关闭连接
• 等待1秒给服务端处理断开事件的时间，确认服务端没有抛出panic，循环正常退出

完整测试代码 #

package main

import (
	"io"
	"log"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

func init() {
	// 测试环境放开Origin校验
	upgrader.CheckOrigin = func(r *http.Request) bool { return true }
	// 测试环境屏蔽业务日志，需要看日志可以注释下面这行
	log.SetOutput(io.Discard)
}

func TestSocketHandler_Echo(t *testing.T) {
	// 启动内存级测试服务
	srv := httptest.NewServer(http.HandlerFunc(socketHandler))
	defer srv.Close()

	// 拼接WS连接地址
	wsAddr := "ws" + strings.TrimPrefix(srv.URL, "http")

	// 建立客户端连接
	client, _, err := websocket.DefaultDialer.Dial(wsAddr, nil)
	if err != nil {
		t.Fatalf("建立WS连接失败: %v", err)
	}
	defer client.Close()

	// 构造测试用例
	testCases := []struct {
		caseName string
		msgType  int
		payload  []byte
	}{
		{"普通文本消息", websocket.TextMessage, []byte("test ws echo")},
		{"二进制消息", websocket.BinaryMessage, []byte{0x10, 0x20, 0x30}},
		{"空消息", websocket.TextMessage, []byte{}},
		{"长文本消息", websocket.TextMessage, []byte(strings.Repeat("a", 1024))},
	}

	for _, tc := range testCases {
		t.Run(tc.caseName, func(t *testing.T) {
			// 发送消息
			if err := client.WriteMessage(tc.msgType, tc.payload); err != nil {
				t.Fatalf("发送消息失败: %v", err)
			}

			// 设置读超时，防止测试卡死
			_ = client.SetReadDeadline(time.Now().Add(2 * time.Second))
			recvType, recvContent, err := client.ReadMessage()
			if err != nil {
				t.Fatalf("读取服务端返回失败: %v", err)
			}

			// 校验返回结果
			if recvType != tc.msgType {
				t.Errorf("消息类型不匹配，期望%d，实际%d", tc.msgType, recvType)
			}
			if string(recvContent) != string(tc.payload) {
				t.Errorf("消息内容不匹配，期望%s，实际%s", tc.payload, recvContent)
			}
		})
	}
}

func TestSocketHandler_UpgradeFail(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(socketHandler))
	defer srv.Close()

	// 发普通HTTP请求，不携带WS握手头
	resp, err := http.Get(srv.URL)
	if err != nil {
		t.Fatalf("发起普通请求失败: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusBadRequest {
		t.Errorf("非法握手时期望状态码400，实际%d", resp.StatusCode)
	}
}

func TestSocketHandler_ClientDisconnect(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(socketHandler))
	defer srv.Close()

	wsAddr := "ws" + strings.TrimPrefix(srv.URL, "http")
	client, _, err := websocket.DefaultDialer.Dial(wsAddr, nil)
	if err != nil {
		t.Fatalf("建立WS连接失败: %v", err)
	}

	// 主动关闭客户端连接
	if err := client.Close(); err != nil {
		t.Fatalf("关闭客户端连接失败: %v", err)
	}

	// 等待服务端处理断开事件，确认无panic
	time.Sleep(1 * time.Second)
}

测试执行 #

在代码目录下直接执行go test -v即可运行所有测试用例，正常情况下所有用例都会在几十毫秒内执行完成。
如果后续给socketHandler增加鉴权、消息校验、广播等逻辑，只需要在现有测试框架基础上增加对应场景的用例即可，不需要改动基础测试结构。

内容的提问来源于stack exchange，提问作者Amit Arora