github.com/GuanceCloud/cliutils@v1.1.21/network/ws/ws_test.go (about)

     1  // Unless explicitly stated otherwise all files in this repository are licensed
     2  // under the MIT License.
     3  // This product includes software developed at Guance Cloud (https://www.guance.com/).
     4  // Copyright 2021-present Guance, Inc.
     5  
     6  package ws
     7  
     8  import (
     9  	"flag"
    10  	"fmt"
    11  	"net"
    12  	"net/http"
    13  	"net/url"
    14  	"sync"
    15  	"testing"
    16  	"time"
    17  
    18  	"github.com/GuanceCloud/cliutils"
    19  	"github.com/gobwas/ws"
    20  	"github.com/gorilla/websocket"
    21  )
    22  
    23  var (
    24  	__wsip     = `0.0.0.0`
    25  	__wsport   = 54321
    26  	__wsupath  = "/wstest"
    27  	__df_wsurl = url.URL{Scheme: "ws", Host: fmt.Sprintf("%s:%d", __wsip, __wsport), Path: __wsupath}
    28  
    29  	__fcliCnt   = flag.Int("cli-cnt", 128, ``)
    30  	__ftestTime = flag.Duration("test-time", time.Minute, ``)
    31  
    32  	__wg = sync.WaitGroup{}
    33  )
    34  
    35  func TestWSServer(t *testing.T) {
    36  	flag.Parse()
    37  
    38  	// ws server
    39  	dfwsurl := fmt.Sprintf("%s:%d", __wsip, __wsport)
    40  	df_srv, err := NewServer(dfwsurl, __wsupath)
    41  	if err != nil {
    42  		t.Fatal(err)
    43  	}
    44  
    45  	// msg-handle callback
    46  	df_srv.MsgHandler = func(s *Server, c net.Conn, data []byte, op ws.OpCode) error {
    47  		SendMsgToClient([]byte(fmt.Sprintf("your are %s", c.RemoteAddr().String())), c)
    48  		return nil
    49  	}
    50  
    51  	// add-cli callback
    52  	df_srv.AddCli = func(w http.ResponseWriter, r *http.Request) {
    53  		conn, _, _, err := ws.UpgradeHTTP(r, w)
    54  		if err != nil {
    55  			l.Error("ws.UpgradeHTTP error: %s", err.Error())
    56  			return
    57  		}
    58  
    59  		id := r.URL.Query().Get("id") // get id from ws request URL
    60  		if id == "" {
    61  			t.Fatal("id miss")
    62  		}
    63  
    64  		if err := df_srv.AddConnection(conn); err != nil {
    65  			l.Error(err)
    66  		}
    67  	}
    68  
    69  	go df_srv.Start()
    70  	time.Sleep(time.Second)
    71  
    72  	type wscli struct {
    73  		id  string
    74  		cli *websocket.Conn
    75  	}
    76  
    77  	ncli := *__fcliCnt
    78  
    79  	// datakit as ws proxy client
    80  	dkclis := []*wscli{}
    81  	for i := 0; i < ncli; i++ {
    82  		cliid := cliutils.XID("id_")
    83  
    84  		dw_wsurl := url.URL{
    85  			Scheme:   "ws",
    86  			Host:     fmt.Sprintf("%s:%d", __wsip, __wsport),
    87  			Path:     __wsupath,
    88  			RawQuery: fmt.Sprintf(`id=%s&version=v1.0.0.0-1234-gabcdef`, cliid),
    89  		}
    90  
    91  		dk_cli, _, err := websocket.DefaultDialer.Dial(dw_wsurl.String(), nil)
    92  		if err != nil {
    93  			t.Fatalf("Failed to connect: %s", err.Error())
    94  		}
    95  		dkclis = append(dkclis, &wscli{id: cliid, cli: dk_cli})
    96  	}
    97  
    98  	__wg.Add(ncli)
    99  	ch := make(chan interface{})
   100  
   101  	// ws-cli send msg to ws-server
   102  	for i := 0; i < ncli; i++ {
   103  		go func(t *testing.T, i int) {
   104  			t.Helper()
   105  
   106  			total := 0
   107  			c := dkclis[i]
   108  			defer __wg.Done()
   109  
   110  			for {
   111  				if err := c.cli.WriteMessage(websocket.TextMessage, []byte(c.id)); err != nil {
   112  					t.Errorf("client write failed: %s", err.Error())
   113  				}
   114  
   115  				total++
   116  				if _, resp, err := c.cli.ReadMessage(); err != nil {
   117  					_ = resp
   118  					t.Log(err)
   119  				} else if total%(ncli/2) == 0 {
   120  					l.Debugf("%s", string(resp))
   121  				}
   122  
   123  				time.Sleep(time.Millisecond)
   124  				select {
   125  				case <-ch:
   126  					c.cli.Close()
   127  					l.Debugf("cli %d exit", i)
   128  					return
   129  				default:
   130  				}
   131  			}
   132  		}(t, i)
   133  	}
   134  
   135  	time.Sleep(*__ftestTime)
   136  	close(ch)
   137  
   138  	df_srv.Stop()
   139  
   140  	__wg.Wait()
   141  }