From 1cc65a821767d14c5e0bd2711717bd22b5421c2f Mon Sep 17 00:00:00 2001 From: Shogo Iwano Date: Wed, 7 Oct 2015 07:01:43 +0900 Subject: [PATCH] Add broadcast methods for a binary message --- melody.go | 25 ++++++++-- melody_test.go | 149 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 171 insertions(+), 3 deletions(-) diff --git a/melody.go b/melody.go index 1a72f24..acf8c45 100644 --- a/melody.go +++ b/melody.go @@ -100,25 +100,44 @@ func (m *Melody) HandleRequest(w http.ResponseWriter, r *http.Request) { go m.disconnectHandler(session) } -// Broadcasts a message to all sessions. +// Broadcasts a text message to all sessions. func (m *Melody) Broadcast(msg []byte) { message := &envelope{t: websocket.TextMessage, msg: msg} m.hub.broadcast <- message } -// Broadcasts a message to all sessions that fn returns true for. +// Broadcasts a text message to all sessions that fn returns true for. func (m *Melody) BroadcastFilter(msg []byte, fn func(*Session) bool) { message := &envelope{t: websocket.TextMessage, msg: msg, filter: fn} m.hub.broadcast <- message } -// Broadcasts a message to all sessions except session s. +// Broadcasts a text message to all sessions except session s. func (m *Melody) BroadcastOthers(msg []byte, s *Session) { m.BroadcastFilter(msg, func(q *Session) bool { return s != q }) } +// Broadcasts a binary message to all sessions. +func (m *Melody) BroadcastBinary(msg []byte) { + message := &envelope{t: websocket.BinaryMessage, msg: msg} + m.hub.broadcast <- message +} + +// Broadcasts a binary message to all sessions that fn returns true for. +func (m *Melody) BroadcastBinaryFilter(msg []byte, fn func(*Session) bool) { + message := &envelope{t: websocket.BinaryMessage, msg: msg, filter: fn} + m.hub.broadcast <- message +} + +// Broadcasts a binary message to all sessions except session s. +func (m *Melody) BroadcastBinaryOthers(msg []byte, s *Session) { + m.BroadcastBinaryFilter(msg, func(q *Session) bool { + return s != q + }) +} + // Closes the melody instance and all connected sessions. func (m *Melody) Close() { m.hub.exit <- true diff --git a/melody_test.go b/melody_test.go index 334bad4..3d5fca5 100644 --- a/melody_test.go +++ b/melody_test.go @@ -1,6 +1,7 @@ package melody import ( + "bytes" "github.com/gorilla/websocket" "net/http" "net/http/httptest" @@ -212,6 +213,56 @@ func TestBroadcast(t *testing.T) { } } +func TestBroadcastBinary(t *testing.T) { + broadcast := NewTestServer() + broadcast.m.HandleMessageBinary(func(session *Session, msg []byte) { + broadcast.m.BroadcastBinary(msg) + }) + server := httptest.NewServer(broadcast) + defer server.Close() + + n := 10 + + fn := func(msg []byte) 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.BinaryMessage, []byte(msg)) + + for i := 0; i < n; i++ { + messageType, ret, err := listeners[i].ReadMessage() + + if err != nil { + t.Error(err) + return false + } + + if messageType != websocket.BinaryMessage { + t.Errorf("message type should be BinaryMessage") + return false + } + + if !bytes.Equal(msg, ret) { + t.Errorf("%v should equal %v", msg, ret) + return false + } + } + + return true + } + + if !fn([]byte{2, 3, 5, 7, 11}) { + t.Errorf("should not be false") + } +} + func TestBroadcastOthers(t *testing.T) { broadcast := NewTestServer() broadcast.m.HandleMessage(func(session *Session, msg []byte) { @@ -259,6 +310,58 @@ func TestBroadcastOthers(t *testing.T) { } } +func TestBroadcastBinaryOthers(t *testing.T) { + broadcast := NewTestServer() + broadcast.m.HandleMessageBinary(func(session *Session, msg []byte) { + broadcast.m.BroadcastBinaryOthers(msg, session) + }) + broadcast.m.Config.PongWait = time.Second + broadcast.m.Config.PingPeriod = time.Second * 9 / 10 + server := httptest.NewServer(broadcast) + defer server.Close() + + n := 10 + + fn := func(msg []byte) 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.BinaryMessage, []byte(msg)) + + for i := 0; i < n; i++ { + messageType, ret, err := listeners[i].ReadMessage() + + if err != nil { + t.Error(err) + return false + } + + if messageType != websocket.BinaryMessage { + t.Errorf("message type should be BinaryMessage") + return false + } + + if !bytes.Equal(msg, ret) { + t.Errorf("%v should equal %v", msg, ret) + return false + } + } + + return true + } + + if !fn([]byte{2, 3, 5, 7, 11}) { + t.Errorf("should not be false") + } +} + func TestPingPong(t *testing.T) { noecho := NewTestServer() noecho.m.Config.PongWait = time.Second @@ -326,6 +429,52 @@ func TestBroadcastFilter(t *testing.T) { } } +func TestBroadcastBinaryFilter(t *testing.T) { + broadcast := NewTestServer() + broadcast.m.HandleMessageBinary(func(session *Session, msg []byte) { + broadcast.m.BroadcastBinaryFilter(msg, func(q *Session) bool { + return session == q + }) + }) + server := httptest.NewServer(broadcast) + defer server.Close() + + fn := func(msg []byte) bool { + conn, err := NewDialer(server.URL) + defer conn.Close() + + if err != nil { + t.Error(err) + return false + } + + conn.WriteMessage(websocket.BinaryMessage, []byte(msg)) + + messageType, ret, err := conn.ReadMessage() + + if err != nil { + t.Error(err) + return false + } + + if messageType != websocket.BinaryMessage { + t.Errorf("message type should be BinaryMessage") + return false + } + + if !bytes.Equal(msg, ret) { + t.Errorf("%v should equal %v", msg, ret) + return false + } + + return true + } + + if !fn([]byte{2, 3, 5, 7, 11}) { + t.Errorf("should not be false") + } +} + func TestStop(t *testing.T) { noecho := NewTestServer() server := httptest.NewServer(noecho)