How to write Unit Test Cases for Gorilla Websocket

Question

I need to write unit test cases for Gorilla Web-socket.

var upgrader = websocket.Upgrader{} // use default options
 
func socketHandler(w http.ResponseWriter, r *http.Request) {
    // Upgrade our raw HTTP connection to a websocket based one
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Print("Error during connection upgradation:", err)
        return
    }
    defer conn.Close()
 
    // The event loop
    for {
        messageType, message, err := conn.ReadMessage()
        if err != nil {
            log.Println("Error during message reading:", err)
            break
        }
        log.Printf("Received: %s", message)
        err = conn.WriteMessage(messageType, message)
        if err != nil {
            log.Println("Error during message writing:", err)
            break
        }
    }
}
 
func home(w http.ResponseWriter, r *http.Request) {
    fmt.Fprintf(w, "Index Page")
}
 
func main() {
    http.HandleFunc("/socket", socketHandler)
    http.HandleFunc("/", home)
    log.Fatal(http.ListenAndServe("localhost:8080", nil))
}

I want to write the Unit Test case for socketHandler function. How should I go about it?

Answer 1

Spin up an httptest.Server with the websocket endpoint. Dial the server with the websocket client. Send and receive messages as needed.

Here's a simple example:

func TestWs(t *testing.T) {
    s := httptest.NewServer(http.HandlerFunc(socketHandler))
    defer s.Close()
    wsURL := "ws" + strings.TrimPrefix(s.URL, "http")
    c, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
    if err != nil {
        t.Error(err)
    }
    defer c.Close()

    // Write a message.
    c.WriteMessage(websocket.TextMessage, []byte("hello"))

    // Expect the server to echo the message back.
    c.SetReadDeadline(time.Now().Add(time.Second * 2))
    mt, msg, err := c.ReadMessage()
    if err != nil {
        t.Error(err)
    }
    if mt != websocket.TextMessage || string(msg) != "hello" {
        t.Errorf("expected text hello, got %d: %s", mt, msg)
    }
}

Run the example on the playground .

To test a client, write a server endpoint that sends as receives messages as needed for a test. Run your client code against the endpoint as above.