github.com/gogf/gf@v1.16.9/net/ghttp/ghttp_unit_websocket_client_test.go (about)

     1  // Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
     2  //
     3  // This Source Code Form is subject to the terms of the MIT License.
     4  // If a copy of the MIT was not distributed with this file,
     5  // You can obtain one at https://github.com/gogf/gf.
     6  
     7  package ghttp_test
     8  
     9  import (
    10  	"fmt"
    11  	"net/http"
    12  	"testing"
    13  	"time"
    14  
    15  	"github.com/gorilla/websocket"
    16  
    17  	"github.com/gogf/gf/frame/g"
    18  	"github.com/gogf/gf/net/ghttp"
    19  	"github.com/gogf/gf/test/gtest"
    20  )
    21  
    22  func Test_WebSocketClient(t *testing.T) {
    23  	p, _ := ports.PopRand()
    24  	s := g.Server(p)
    25  	s.BindHandler("/ws", func(r *ghttp.Request) {
    26  		ws, err := r.WebSocket()
    27  		if err != nil {
    28  			r.Exit()
    29  		}
    30  		for {
    31  			msgType, msg, err := ws.ReadMessage()
    32  			if err != nil {
    33  				return
    34  			}
    35  			if err = ws.WriteMessage(msgType, msg); err != nil {
    36  				return
    37  			}
    38  		}
    39  	})
    40  	s.SetPort(p)
    41  	s.SetDumpRouterMap(false)
    42  	s.Start()
    43  	defer s.Shutdown()
    44  
    45  	time.Sleep(100 * time.Millisecond)
    46  	gtest.C(t, func(t *gtest.T) {
    47  		client := ghttp.NewWebSocketClient()
    48  		client.Proxy = http.ProxyFromEnvironment
    49  		client.HandshakeTimeout = time.Minute
    50  
    51  		conn, _, err := client.Dial(fmt.Sprintf("ws://127.0.0.1:%d/ws", p), nil)
    52  		t.Assert(err, nil)
    53  		defer conn.Close()
    54  
    55  		msg := []byte("hello")
    56  		err = conn.WriteMessage(websocket.TextMessage, msg)
    57  		t.Assert(err, nil)
    58  
    59  		mt, data, err := conn.ReadMessage()
    60  		t.Assert(err, nil)
    61  		t.Assert(mt, websocket.TextMessage)
    62  		t.Assert(data, msg)
    63  	})
    64  }