Use Upgrader in client/server test.

This commit is contained in:
Gary Burd 2014-04-20 12:08:24 -07:00
parent df94ef6daa
commit bbe4cde693
1 changed files with 50 additions and 55 deletions

View File

@ -17,6 +17,42 @@ import (
"time" "time"
) )
func sendRecv(t *testing.T, ws *Conn) {
const message = "Hello World!"
if err := ws.SetWriteDeadline(time.Now().Add(time.Second)); err != nil {
t.Fatalf("SetWriteDeadline: %v", err)
}
if err := ws.WriteMessage(TextMessage, []byte(message)); err != nil {
t.Fatalf("WriteMessage: %v", err)
}
if err := ws.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
t.Fatalf("SetReadDeadline: %v", err)
}
_, p, err := ws.ReadMessage()
if err != nil {
t.Fatalf("ReadMessage: %v", err)
}
if string(p) != message {
t.Fatalf("message=%s, want %s", p, message)
}
}
func httpToWs(u string) string {
return "ws" + u[len("http"):]
}
var handshakeUpgrader = &Upgrader{
Subprotocols: []string{"p0", "p1"},
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
var handshakeDialer = &Dialer{
Subprotocols: []string{"p1", "p2"},
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
type handshakeHandler struct { type handshakeHandler struct {
*testing.T *testing.T
} }
@ -27,33 +63,21 @@ func (t handshakeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
t.Logf("method = %s, want GET", r.Method) t.Logf("method = %s, want GET", r.Method)
return return
} }
if origin := r.Header.Get("Origin"); origin != "http://"+r.Host {
http.Error(w, "Origin not allowed", 403)
t.Logf("Origin = %s, want %s", origin, r.Host)
return
}
subprotos := Subprotocols(r) subprotos := Subprotocols(r)
if !reflect.DeepEqual(subprotos, handshakeDialer.Subprotocols) { if !reflect.DeepEqual(subprotos, handshakeDialer.Subprotocols) {
http.Error(w, "bad protocol", 400) http.Error(w, "bad protocol", 400)
t.Logf("Subprotocols = %v, want %v", subprotos, handshakeDialer.Subprotocols) t.Logf("Subprotocols = %v, want %v", subprotos, handshakeDialer.Subprotocols)
return return
} }
ws, err := Upgrade(w, r, http.Header{ ws, err := handshakeUpgrader.Upgrade(w, r, http.Header{"Set-Cookie": {"sessionID=1234"}})
"Set-Cookie": {"sessionID=1234"}, if err != nil {
"Sec-Websocket-Protocol": {subprotos[0]},
}, 1024, 1024)
if _, ok := err.(HandshakeError); ok {
t.Logf("bad handshake: %v", err)
http.Error(w, "Not a websocket handshake", 400)
return
} else if err != nil {
t.Logf("upgrade error: %v", err) t.Logf("upgrade error: %v", err)
return return
} }
defer ws.Close() defer ws.Close()
if ws.Subprotocol() != subprotos[0] { if ws.Subprotocol() != "p1" {
t.Logf("ws.Subprotocol() = %s, want %s", ws.Subprotocol(), subprotos[0]) t.Logf("ws.Subprotocol() = %s, want p1", ws.Subprotocol())
return return
} }
@ -81,12 +105,6 @@ func (t handshakeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
} }
} }
var handshakeDialer = &Dialer{
Subprotocols: []string{"p1", "p2"},
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
func TestHandshake(t *testing.T) { func TestHandshake(t *testing.T) {
s := httptest.NewServer(handshakeHandler{t}) s := httptest.NewServer(handshakeHandler{t})
defer s.Close() defer s.Close()
@ -106,10 +124,9 @@ func TestHandshake(t *testing.T) {
t.Error("Set-Cookie not received from the server.") t.Error("Set-Cookie not received from the server.")
} }
if ws.Subprotocol() != handshakeDialer.Subprotocols[0] { if ws.Subprotocol() != "p1" {
t.Errorf("ws.Subprotocol() = %s, want %s", ws.Subprotocol(), handshakeDialer.Subprotocols[0]) t.Errorf("ws.Subprotocol() = %s, want p1", ws.Subprotocol())
} }
sendRecv(t, ws) sendRecv(t, ws)
} }
@ -117,13 +134,15 @@ type dialHandler struct {
*testing.T *testing.T
} }
var dialUpgrader = &Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
CheckOrigin: func(r *http.Request) bool { return true },
}
func (t dialHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { func (t dialHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
ws, err := Upgrade(w, r, nil, 1024, 1024) ws, err := dialUpgrader.Upgrade(w, r, nil)
if _, ok := err.(HandshakeError); ok { if err != nil {
t.Logf("bad handshake: %v", err)
http.Error(w, "Not a websocket handshake", 400)
return
} else if err != nil {
t.Logf("upgrade error: %v", err) t.Logf("upgrade error: %v", err)
return return
} }
@ -143,30 +162,6 @@ func (t dialHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
} }
} }
func sendRecv(t *testing.T, ws *Conn) {
const message = "Hello World!"
if err := ws.SetWriteDeadline(time.Now().Add(time.Second)); err != nil {
t.Fatalf("SetWriteDeadline: %v", err)
}
if err := ws.WriteMessage(TextMessage, []byte(message)); err != nil {
t.Fatalf("WriteMessage: %v", err)
}
if err := ws.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
t.Fatalf("SetReadDeadline: %v", err)
}
_, p, err := ws.ReadMessage()
if err != nil {
t.Fatalf("ReadMessage: %v", err)
}
if string(p) != message {
t.Fatalf("message=%s, want %s", p, message)
}
}
func httpToWs(u string) string {
return "ws" + u[len("http"):]
}
func TestDial(t *testing.T) { func TestDial(t *testing.T) {
s := httptest.NewServer(dialHandler{t}) s := httptest.NewServer(dialHandler{t})
defer s.Close() defer s.Close()