在golang中为ws创建单元测试

package main_test

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

	"github.com/gorilla/websocket"
)

func TestWebSocketServer(t *testing.T) {
	// 创建测试服务器
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// 升级连接为WebSocket
		upgrader := websocket.Upgrader{}
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			t.Fatalf("Failed to upgrade connection to WebSocket: %v", err)
		}

		// 读取客户端发送的消息
		_, msg, err := conn.ReadMessage()
		if err != nil {
			t.Fatalf("Failed to read message from client: %v", err)
		}

		// 发送响应消息给客户端
		err = conn.WriteMessage(websocket.TextMessage, msg)
		if err != nil {
			t.Fatalf("Failed to write message to client: %v", err)
		}
	}))
	defer server.Close()

	// 连接到测试服务器
	url := "ws" + server.URL[4:]
	conn, _, err := websocket.DefaultDialer.Dial(url, nil)
	if err != nil {
		t.Fatalf("Failed to connect to server: %v", err)
	}
	defer conn.Close()

	// 发送消息到服务器
	msg := []byte("Hello, server!")
	err = conn.WriteMessage(websocket.TextMessage, msg)
	if err != nil {
		t.Fatalf("Failed to send message to server: %v", err)
	}

	// 接收服务器的响应消息
	_, resp, err := conn.ReadMessage()
	if err != nil {
		t.Fatalf("Failed to read message from server: %v", err)
	}

	// 验证服务器的响应消息是否正确
	if string(resp) != string(msg) {
		t.Fatalf("Unexpected response from server: %s", string(resp))
	}
}