From 59e6272ba4fa10d7bb5625a65e98b51186ab0b6d Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Tue, 26 May 2015 23:28:45 +0900
Subject: [PATCH] integration: Add WebSocket upgrade test

---
 README.rst                              |  1 +
 integration-tests/Makefile.am           |  1 +
 integration-tests/nghttpx_http1_test.go | 22 +++++++++++++
 integration-tests/server_tester.go      | 42 ++++++++++++++++++++++++-
 4 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/README.rst b/README.rst
index e8014b5f..1092ebbf 100644
--- a/README.rst
+++ b/README.rst
@@ -191,6 +191,7 @@ its testing framework.  We depend on the following libraries:
 * https://github.com/bradfitz/http2
 * https://github.com/tatsuhiro-t/go-nghttp2
 * https://github.com/tatsuhiro-t/spdy
+* golang.org/x/net/websocket
 
 To download the above packages, after settings ``GOPATH``, run the
 following command under ``integration-tests`` directory::
diff --git a/integration-tests/Makefile.am b/integration-tests/Makefile.am
index 0ed3273f..2ab15608 100644
--- a/integration-tests/Makefile.am
+++ b/integration-tests/Makefile.am
@@ -38,6 +38,7 @@ itprep:
 	go get -d -v github.com/bradfitz/http2
 	go get -d -v github.com/tatsuhiro-t/go-nghttp2
 	go get -d -v github.com/tatsuhiro-t/spdy
+	go get -d -v golang.org/x/net/websocket
 
 it:
 	sh setenv go test -v
diff --git a/integration-tests/nghttpx_http1_test.go b/integration-tests/nghttpx_http1_test.go
index c8d58709..6d5c28ad 100644
--- a/integration-tests/nghttpx_http1_test.go
+++ b/integration-tests/nghttpx_http1_test.go
@@ -2,8 +2,10 @@ package nghttp2
 
 import (
 	"bufio"
+	"bytes"
 	"fmt"
 	"github.com/bradfitz/http2/hpack"
+	"golang.org/x/net/websocket"
 	"io"
 	"net/http"
 	"syscall"
@@ -312,6 +314,26 @@ func TestH1H1HeaderFields(t *testing.T) {
 	}
 }
 
+// TestH1H1Websocket tests that HTTP Upgrade to WebSocket works.
+func TestH1H1Websocket(t *testing.T) {
+	st := newServerTesterHandler(nil, t, websocket.Handler(func(ws *websocket.Conn) {
+		io.Copy(ws, ws)
+	}))
+	defer st.Close()
+
+	content := []byte("hello world")
+	res, err := st.websocket(requestParam{
+		name: "TestH1H1Websocket",
+		body: content,
+	})
+	if err != nil {
+		t.Fatalf("Error st.websocket() = %v", err)
+	}
+	if got, want := res.body, content; !bytes.Equal(got, want) {
+		t.Errorf("echo: %q; want %q", got, want)
+	}
+}
+
 // TestH1H2ConnectFailure tests that server handles the situation that
 // connection attempt to HTTP/2 backend failed.
 func TestH1H2ConnectFailure(t *testing.T) {
diff --git a/integration-tests/server_tester.go b/integration-tests/server_tester.go
index 43407956..689ffa89 100644
--- a/integration-tests/server_tester.go
+++ b/integration-tests/server_tester.go
@@ -10,6 +10,7 @@ import (
 	"github.com/bradfitz/http2/hpack"
 	"github.com/tatsuhiro-t/go-nghttp2"
 	"github.com/tatsuhiro-t/spdy"
+	"golang.org/x/net/websocket"
 	"io"
 	"io/ioutil"
 	"net"
@@ -66,6 +67,10 @@ func newServerTester(args []string, t *testing.T, handler http.HandlerFunc) *ser
 	return newServerTesterInternal(args, t, handler, false, nil)
 }
 
+func newServerTesterHandler(args []string, t *testing.T, handler http.Handler) *serverTester {
+	return newServerTesterInternal(args, t, handler, false, nil)
+}
+
 // newServerTester creates test context for TLS frontend connection.
 func newServerTesterTLS(args []string, t *testing.T, handler http.HandlerFunc) *serverTester {
 	return newServerTesterInternal(args, t, handler, true, nil)
@@ -79,7 +84,7 @@ func newServerTesterTLSConfig(args []string, t *testing.T, handler http.HandlerF
 
 // newServerTesterInternal creates test context.  If frontendTLS is
 // true, set up TLS frontend connection.
-func newServerTesterInternal(args []string, t *testing.T, handler http.HandlerFunc, frontendTLS bool, clientConfig *tls.Config) *serverTester {
+func newServerTesterInternal(args []string, t *testing.T, handler http.Handler, frontendTLS bool, clientConfig *tls.Config) *serverTester {
 	ts := httptest.NewUnstartedServer(handler)
 
 	backendTLS := false
@@ -279,6 +284,41 @@ func (cbr *chunkedBodyReader) Read(p []byte) (n int, err error) {
 	return cbr.body.Read(p)
 }
 
+func (st *serverTester) websocket(rp requestParam) (*serverResponse, error) {
+	urlstring := st.url + "/echo"
+
+	config, err := websocket.NewConfig(urlstring, st.url)
+	if err != nil {
+		st.t.Fatalf("websocket.NewConfig(%q, %q) returned error: %v", urlstring, st.url, err)
+	}
+
+	config.Header.Add("Test-Case", rp.name)
+	for _, h := range rp.header {
+		config.Header.Add(h.Name, h.Value)
+	}
+
+	ws, err := websocket.NewClient(config, st.conn)
+	if err != nil {
+		st.t.Fatalf("Error creating websocket client: %v", err)
+	}
+
+	if _, err := ws.Write(rp.body); err != nil {
+		st.t.Fatalf("ws.Write() returned error: %v", err)
+	}
+
+	msg := make([]byte, 1024)
+	var n int
+	if n, err = ws.Read(msg); err != nil {
+		st.t.Fatalf("ws.Read() returned error: %v", err)
+	}
+
+	res := &serverResponse{
+		body: msg[:n],
+	}
+
+	return res, nil
+}
+
 func (st *serverTester) http1(rp requestParam) (*serverResponse, error) {
 	method := "GET"
 	if rp.method != "" {