websocket 增加多分组 fork https://github.com/olahol/melody
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

239 lines
4.4 KiB

package melody
import (
"github.com/gorilla/websocket"
"net/http"
"net/http/httptest"
"strings"
"testing"
"testing/quick"
"time"
)
type TestServer struct {
m *Melody
}
func NewTestServerHandler(handler handleMessageFunc) *TestServer {
m := New()
m.HandleMessage(handler)
return &TestServer{
m: m,
}
}
func NewTestServer() *TestServer {
m := New()
return &TestServer{
m: m,
}
}
func (s *TestServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
s.m.HandleRequest(w, r)
}
func NewDialer(url string) (*websocket.Conn, error) {
dialer := &websocket.Dialer{}
conn, _, err := dialer.Dial(strings.Replace(url, "http", "ws", 1), nil)
return conn, err
}
func TestEcho(t *testing.T) {
echo := NewTestServerHandler(func(session *Session, msg []byte) {
session.Write(msg)
})
server := httptest.NewServer(echo)
defer server.Close()
fn := func(msg string) bool {
conn, err := NewDialer(server.URL)
defer conn.Close()
if err != nil {
t.Error(err)
return false
}
conn.WriteMessage(websocket.TextMessage, []byte(msg))
_, ret, err := conn.ReadMessage()
if err != nil {
t.Error(err)
return false
}
if msg != string(ret) {
t.Errorf("%s should equal %s", msg, string(ret))
return false
}
return true
}
if err := quick.Check(fn, nil); err != nil {
t.Error(err)
}
}
func TestBroadcast(t *testing.T) {
broadcast := NewTestServer()
broadcast.m.HandleMessage(func(session *Session, msg []byte) {
broadcast.m.Broadcast(msg)
})
server := httptest.NewServer(broadcast)
defer server.Close()
n := 10
fn := func(msg string) bool {
conn, _ := NewDialer(server.URL)
defer conn.Close()
listeners := make([]*websocket.Conn, n)
for i := 0; i < n; i++ {
listener, _ := NewDialer(server.URL)
listeners[i] = listener
defer listeners[i].Close()
}
conn.WriteMessage(websocket.TextMessage, []byte(msg))
for i := 0; i < n; i++ {
_, ret, err := listeners[i].ReadMessage()
if err != nil {
t.Error(err)
return false
}
if msg != string(ret) {
t.Errorf("%s should equal %s", msg, string(ret))
return false
}
}
return true
}
if err := quick.Check(fn, nil); err != nil {
t.Error(err)
}
}
func TestPingPong(t *testing.T) {
noecho := NewTestServer()
noecho.m.Config.PongWait = time.Second
noecho.m.Config.PingPeriod = time.Second * 9 / 10
server := httptest.NewServer(noecho)
defer server.Close()
conn, err := NewDialer(server.URL)
conn.SetPingHandler(func(string) error {
return nil
})
defer conn.Close()
if err != nil {
t.Error(err)
}
conn.WriteMessage(websocket.TextMessage, []byte("test"))
_, _, err = conn.ReadMessage()
if err == nil {
t.Error("there should be an error")
}
}
/*
func TestBroadcastFilter(t *testing.T) {
echo := NewTestServer()
echo.m.HandleMessage(func(session *Session, msg []byte) {
echo.m.BroadcastFilter(func(s *Session) bool {
//return s == session
return false
}, msg)
})
server := httptest.NewServer(echo)
defer server.Close()
fn := func(msg string) bool {
conn, err := NewDialer(server.URL)
conn.SetPingHandler(func(string) error {
return nil
})
defer conn.Close()
if err != nil {
t.Error(err)
return false
}
conn.WriteMessage(websocket.TextMessage, []byte(msg))
_, ret, err := conn.ReadMessage()
if err != nil {
t.Error(err)
return false
}
if msg != string(ret) {
t.Errorf("%s should equal %s", msg, string(ret))
return false
}
return true
}
if err := quick.Check(fn, nil); err != nil {
t.Error(err)
}
}
*/
func BenchmarkEcho(b *testing.B) {
echo := NewTestServerHandler(func(session *Session, msg []byte) {
session.Write(msg)
})
server := httptest.NewServer(echo)
defer server.Close()
conn, _ := NewDialer(server.URL)
defer conn.Close()
for i := 0; i < b.N; i++ {
conn.WriteMessage(websocket.TextMessage, []byte("test"))
conn.ReadMessage()
}
}
func BenchmarkBroadcast(b *testing.B) {
broadcast := NewTestServer()
broadcast.m.HandleMessage(func(session *Session, msg []byte) {
broadcast.m.Broadcast(msg)
})
server := httptest.NewServer(broadcast)
defer server.Close()
conn, _ := NewDialer(server.URL)
defer conn.Close()
n := 10
listeners := make([]*websocket.Conn, n)
for i := 0; i < n; i++ {
listener, _ := NewDialer(server.URL)
listeners[i] = listener
defer listeners[i].Close()
}
for i := 0; i < b.N; i++ {
conn.WriteMessage(websocket.TextMessage, []byte("test"))
for i := 0; i < n; i++ {
listeners[i].ReadMessage()
}
}
}