diff --git a/melody.go b/melody.go index 21a1da7..3d0ce8a 100644 --- a/melody.go +++ b/melody.go @@ -12,7 +12,7 @@ type filterFunc func(*Session) bool type Melody struct { Config *Config - upgrader *websocket.Upgrader + Upgrader *websocket.Upgrader messageHandler handleMessageFunc errorHandler handleErrorFunc connectHandler handleSessionFunc @@ -33,7 +33,7 @@ func New() *Melody { return &Melody{ Config: newConfig(), - upgrader: upgrader, + Upgrader: upgrader, messageHandler: func(*Session, []byte) {}, errorHandler: func(*Session, error) {}, connectHandler: func(*Session) {}, @@ -64,9 +64,10 @@ func (m *Melody) HandleError(fn func(*Session, error)) { // Handles a http request and upgrades it to a websocket. func (m *Melody) HandleRequest(w http.ResponseWriter, r *http.Request) { - conn, err := m.upgrader.Upgrade(w, r, nil) + conn, err := m.Upgrader.Upgrade(w, r, nil) if err != nil { + m.errorHandler(nil, err) return } diff --git a/melody_test.go b/melody_test.go index 0b6557a..dcbfe97 100644 --- a/melody_test.go +++ b/melody_test.go @@ -77,6 +77,33 @@ func TestEcho(t *testing.T) { } } +func TestUpgrader(t *testing.T) { + broadcast := NewTestServer() + broadcast.m.HandleMessage(func(session *Session, msg []byte) { + session.Write(msg) + }) + server := httptest.NewServer(broadcast) + defer server.Close() + + broadcast.m.Upgrader = &websocket.Upgrader{ + ReadBufferSize: 1024, + WriteBufferSize: 1024, + CheckOrigin: func(r *http.Request) bool { return false }, + } + + broadcast.m.HandleError(func(session *Session, err error) { + if err == nil || err.Error() != "websocket: origin not allowed" { + t.Error("there should be a origin error") + } + }) + + _, err := NewDialer(server.URL) + + if err == nil || err.Error() != "websocket: bad handshake" { + t.Error("there should be a badhandshake error") + } +} + func TestBroadcast(t *testing.T) { broadcast := NewTestServer() broadcast.m.HandleMessage(func(session *Session, msg []byte) {