Added test from Echo.WebSocket()

Signed-off-by: Vishal Rana <vr@labstack.com>

echo.go

@@ -120,9 +120,9 @@ var (
 	RendererNotRegistered = errors.New("echo ⇒ renderer not registered")
 )
 
-func NewHTTPError(code int, msgs ...string) *HTTPError {
+func NewHTTPError(code int, msg ...string) *HTTPError {
 	he := &HTTPError{Code: code, Message: http.StatusText(code)}
-	for _, m := range msgs {
+	for _, m := range msg {
 		he.Message = m
 	}
 	return he
@@ -272,16 +272,16 @@ func (e *Echo) Trace(path string, h Handler) {
 
 // WebSocket adds a WebSocket route > handler to the router.
 func (e *Echo) WebSocket(path string, h HandlerFunc) {
-	e.Get(path, func(c *Context) *HTTPError {
+	e.Get(path, func(c *Context) (err error) {
 		wss := websocket.Server{
 			Handler: func(ws *websocket.Conn) {
 				c.Socket = ws
 				c.Response.status = http.StatusSwitchingProtocols
-				h(c)
+				err = h(c)
 			},
 		}
 		wss.ServeHTTP(c.Response.writer, c.Request)
-		return nil
+		return err
 	})
 }
 

echo_test.go

@@ -2,9 +2,12 @@ package echo
 
 import (
 	"bytes"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+
+	"golang.org/x/net/websocket"
 )
 
 type (
@@ -246,6 +249,31 @@ func TestEchoMethod(t *testing.T) {
 	e.Trace("/", h)
 }
 
+func TestWebSocket(t *testing.T) {
+	e := New()
+	e.WebSocket("/ws", func(c *Context) error {
+		c.Socket.Write([]byte("test"))
+		return nil
+	})
+	srv := httptest.NewServer(e)
+	defer srv.Close()
+	addr := srv.Listener.Addr().String()
+	origin := "http://localhost"
+	url := fmt.Sprintf("ws://%s/ws", addr)
+	ws, err := websocket.Dial(url, "", origin)
+	if err != nil {
+		t.Fatal(err)
+	}
+	ws.Write([]byte("test"))
+	defer ws.Close()
+	buf := new(bytes.Buffer)
+	buf.ReadFrom(ws)
+	s := buf.String()
+	if s != "test" {
+		t.Errorf("expected `test`, got %s.", s)
+	}
+}
+
 func TestEchoURL(t *testing.T) {
 	e := New()
 	static := func(*Context) error { return nil }

middleware/compress_test.go

@@ -2,11 +2,12 @@ package middleware
 
 import (
 	"compress/gzip"
-	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 
+	"bytes"
+
 	"github.com/labstack/echo"
 )
 
@@ -39,13 +40,11 @@ func TestGzip(t *testing.T) {
 	r, err := gzip.NewReader(w.Body)
 	defer r.Close()
 	if err != nil {
-		t.Error(err)
+		t.Fatal(err)
 	}
-	b, err := ioutil.ReadAll(r)
-	if err != nil {
-		t.Error(err)
-	}
-	s = string(b)
+	buf := new(bytes.Buffer)
+	buf.ReadFrom(r)
+	s = buf.String()
 	if s != "test" {
 		t.Errorf("expected body `test`, got %s.", s)
 	}