Golang基于Gorilla Toolkit的WebSocket客户端与服务端测试用例编写方法

测试Gorilla实现的WebSocket逻辑不需要提前启动独立服务、也不需要依赖外部网络环境，用Go标准库的net/http/httptest配合Gorilla本身的Dialer、Upgrader就能完成全链路闭环测试，覆盖所有核心场景。

服务端测试 #

服务端测试的核心思路是用测试框架启动临时HTTP服务，挂载你实际写的WebSocket业务逻辑，再用模拟客户端发起真实WebSocket请求做校验，不需要修改原有业务代码。

• 核心测试场景覆盖：
  • 连接升级校验：验证合法的WebSocket请求能正常升级，非法请求（不带升级头、跨域非法等）能被正确拦截
  • 消息收发逻辑：验证服务端能正确解析客户端发送的消息、返回符合预期的响应（比如echo逻辑就要校验返回内容和发送内容完全一致）
  • 连接关闭处理：验证客户端主动发送关闭帧、异常断开时，服务端能正常释放连接，不会出现goroutine泄漏
  • 边界场景：验证超大消息、二进制消息、ping/pong心跳处理逻辑符合预期
• 参考测试代码：

import (
    "bytes"
    "net/http"
    "net/http/httptest"
    "net/url"
    "testing"
    "github.com/gorilla/websocket"
)

// 替换成你实际实现的echo服务端handler
func echoHandler(w http.ResponseWriter, r *http.Request) {
    upgrader := websocket.Upgrader{}
    c, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        return
    }
    defer c.Close()
    for {
        mt, message, err := c.ReadMessage()
        if err != nil {
            break
        }
        err = c.WriteMessage(mt, message)
        if err != nil {
            break
        }
    }
}

func TestEchoServer(t *testing.T) {
    // 启动临时测试服务，挂载业务handler
    srv := httptest.NewServer(http.HandlerFunc(echoHandler))
    defer srv.Close()

    // 拼接ws协议的连接地址
    srvURL, _ := url.Parse(srv.URL)
    srvURL.Scheme = "ws"
    srvURL.Path = "/echo"

    // 模拟客户端拨号
    cliConn, _, err := websocket.DefaultDialer.Dial(srvURL.String(), nil)
    if err != nil {
        t.Fatalf("拨号失败: %v", err)
    }
    defer cliConn.Close()

    // 测试正常消息收发
    testCases := []struct {
        name string
        msg  []byte
    }{
        {"普通文本消息", []byte("hello websocket")},
        {"二进制消息", []byte{0x00, 0x01, 0x02}},
        {"空消息", []byte{}},
    }
    for _, tc := range testCases {
        t.Run(tc.name, func(t *testing.T) {
            err = cliConn.WriteMessage(websocket.TextMessage, tc.msg)
            if err != nil {
                t.Fatalf("发送消息失败: %v", err)
            }
            mt, recv, err := cliConn.ReadMessage()
            if err != nil {
                t.Fatalf("接收消息失败: %v", err)
            }
            if mt != websocket.TextMessage {
                t.Errorf("消息类型错误，期望文本类型，实际为%d", mt)
            }
            if !bytes.Equal(recv, tc.msg) {
                t.Errorf("返回内容不匹配，期望%s，实际%s", tc.msg, recv)
            }
        })
    }

    // 测试正常关闭流程
    err = cliConn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
    if err != nil {
        t.Fatalf("发送关闭帧失败: %v", err)
    }
    _, _, err = cliConn.ReadMessage()
    if !websocket.IsCloseError(err, websocket.CloseNormalClosure) {
        t.Errorf("未收到正常关闭响应，错误: %v", err)
    }
}

客户端测试 #

客户端测试的核心思路是在测试逻辑里启动临时WebSocket服务端，模拟真实服务的响应行为，替换客户端默认连接地址后启动客户端逻辑，校验客户端的行为符合预期。

注意：不要直接在测试里调用你写的main函数，建议把main里的核心逻辑抽成独立的可传参、可控退出的函数（比如runClient(addr string, interrupt <-chan os.Signal) error），把原来log.Fatal这类直接退出进程的逻辑改成返回error，方便测试调用。

• 核心测试场景覆盖：
  • 拨号逻辑：客户端能正常和服务端建立连接，地址、请求头符合预期
  • 定时发送逻辑：客户端会按设定的1秒间隔发送时间格式的消息
  • 消息接收逻辑：收到服务端推送的消息后能正确处理
  • 信号处理逻辑：收到中断信号后，客户端能主动发送正常关闭帧，优雅退出不报错
• 参考测试代码：

import (
    "net/http"
    "net/http/httptest"
    "os"
    "testing"
    "time"
    "github.com/gorilla/websocket"
)

// 替换成你抽离出来的客户端核心运行逻辑
func runClient(addr string, interrupt <-chan os.Signal) error {
    // 把原main里的逻辑搬过来，去掉flag解析、log.Fatal部分，参数从外部传入
    // ... 原有业务逻辑
    return nil
}

func TestClient(t *testing.T) {
    msgRecv := make(chan []byte, 1)
    closeRecv := make(chan struct{})

    // 启动测试用WebSocket服务端
    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        upgrader := websocket.Upgrader{}
        srvConn, err := upgrader.Upgrade(w, r, nil)
        if err != nil {
            t.Fatalf("测试服务端升级连接失败: %v", err)
        }
        defer srvConn.Close()

        // 读取客户端发的第一条定时消息
        _, msg, err := srvConn.ReadMessage()
        if err != nil {
            if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
                close(closeRecv)
                return
            }
            t.Errorf("读客户端消息失败: %v", err)
            return
        }
        msgRecv <- msg

        // 给客户端推一条测试消息
        _ = srvConn.WriteMessage(websocket.TextMessage, []byte("server push"))
    }))
    defer srv.Close()

    // 替换客户端连接地址为测试服务地址
    testAddr := srv.Listener.Addr().String()
    interruptSig := make(chan os.Signal, 1)

    // 启动客户端
    clientDone := make(chan error, 1)
    go func() {
        clientDone <- runClient(testAddr, interruptSig)
    }()

    // 校验客户端定时发消息逻辑
    select {
    case sentMsg := <-msgRecv:
        if len(sentMsg) == 0 {
            t.Error("客户端发送的消息为空")
        }
        t.Logf("客户端发送消息: %s", sentMsg)
    case <-time.After(2 * time.Second):
        t.Fatal("等待客户端发送消息超时")
    }

    // 模拟发送中断信号，测试优雅退出
    interruptSig <- os.Interrupt
    select {
    case <-closeRecv:
        t.Log("客户端正常发送关闭帧，连接断开")
    case err := <-clientDone:
        if err != nil {
            t.Errorf("客户端退出报错: %v", err)
        }
    case <-time.After(2 * time.Second):
        t.Fatal("等待客户端退出超时")
    }
}

测试注意事项 #

• 所有测试启动的临时服务、建立的连接都要加defer做资源释放，避免测试跑完残留goroutine或者端口占用
• 涉及时间等待的场景不要写固定的长时间sleep，用channel做同步，超时时间设置为2-3倍的逻辑执行时间即可，提升测试运行速度
• 并发场景测试可以直接启动多个客户端协程同时拨号，用sync.WaitGroup等所有协程执行完再做结果校验即可
• 所有测试逻辑要做到本地闭环，不要依赖外部公共服务，保证测试在任何环境下都能重复稳定运行

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