Refactor WS test server to optional handlers

grafana · Dec 23, 2021 · 6a12e2a · 6a12e2a
1 parent ba3d347
commit 6a12e2a
Show file tree

Hide file tree

Showing 3 changed files with 136 additions and 134 deletions.
diff --git a/common/connection_test.go b/common/connection_test.go
@@ -38,7 +38,7 @@ import (
 )
 
 func TestConnection(t *testing.T) {
-	server := ws.NewServerWithEcho(t)
+	server := ws.NewServer(t, ws.WithEchoHandler("/echo"))
 
 	t.Run("connect", func(t *testing.T) {
 		ctx := context.Background()
@@ -52,7 +52,7 @@ func TestConnection(t *testing.T) {
 }
 
 func TestConnectionClosureAbnormal(t *testing.T) {
-	server := ws.NewServerWithClosureAbnormal(t)
+	server := ws.NewServer(t, ws.WithClosureAbnormalHandler("/closure-abnormal"))
 
 	t.Run("closure abnormal", func(t *testing.T) {
 		ctx := context.Background()
@@ -69,7 +69,7 @@ func TestConnectionClosureAbnormal(t *testing.T) {
 }
 
 func TestConnectionSendRecv(t *testing.T) {
-	server := ws.NewServerWithCDPHandler(t, ws.CDPDefaultHandler, nil)
+	server := ws.NewServer(t, ws.WithCDPHandler("/cdp", ws.CDPDefaultHandler, nil))
 
 	t.Run("send command with empty reply", func(t *testing.T) {
 		ctx := context.Background()
@@ -108,19 +108,19 @@ func TestConnectionCreateSession(t *testing.T) {
 					writeCh <- cdproto.Message{
 						Method: cdproto.EventTargetAttachedToTarget,
 						Params: easyjson.RawMessage([]byte(`
-                            {
-                                "sessionId": "0123456789",
-                                "targetInfo": {
-                                    "targetId": "abcdef0123456789",
-                                    "type": "page",
-                                    "title": "",
-                                    "url": "about:blank",
-                                    "attached": true,
-                                    "browserContextId": "0123456789876543210"
-                                },
-                                "waitingForDebugger": false
-                            }
-                        `)),
+						{
+							"sessionId": "0123456789",
+							"targetInfo": {
+								"targetId": "abcdef0123456789",
+								"type": "page",
+								"title": "",
+								"url": "about:blank",
+								"attached": true,
+								"browserContextId": "0123456789876543210"
+							},
+							"waitingForDebugger": false
+						}
+						`)),
 					}
 					writeCh <- cdproto.Message{
 						ID:        msg.ID,
@@ -132,7 +132,7 @@ func TestConnectionCreateSession(t *testing.T) {
 		}
 	}
 
-	server := ws.NewServerWithCDPHandler(t, handler, &cmdsReceived)
+	server := ws.NewServer(t, ws.WithCDPHandler("/cdp", handler, &cmdsReceived))
 
 	t.Run("create session for target", func(t *testing.T) {
 		ctx := context.Background()

diff --git a/common/session_test.go b/common/session_test.go
@@ -72,7 +72,7 @@ func TestSessionCreateSession(t *testing.T) {
 		}
 	}
 
-	server := ws.NewServerWithCDPHandler(t, handler, &cmdsReceived)
+	server := ws.NewServer(t, ws.WithCDPHandler("/cdp", handler, &cmdsReceived))
 
 	t.Run("send and recv session commands", func(t *testing.T) {
 		ctx := context.Background()

diff --git a/tests/ws/server.go b/tests/ws/server.go
@@ -71,40 +71,22 @@ var (
 	TargetAttachedToTargetResult = fmt.Sprintf(`{"sessionId":"%s"}`, DummyCDPSessionID)
 )
 
-// NewServerWithCDPHandler creates a WS test server with a custom CDP handler function
-func NewServerWithCDPHandler(
-	t testing.TB,
-	fn func(conn *websocket.Conn, msg *cdproto.Message, writeCh chan cdproto.Message, done chan struct{}),
-	cmdsReceived *[]cdproto.MethodType) *Server {
-	return NewServer(t, "/cdp", getWebsocketHandlerCDP(fn, cmdsReceived))
-}
-
 // Server can be used as a test alternative to a real CDP compatible browser.
 type Server struct {
+	t             testing.TB
 	Mux           *http.ServeMux
 	ServerHTTP    *httptest.Server
 	Dialer        *k6netext.Dialer
 	HTTPTransport *http.Transport
 	Context       context.Context
 }
 
-// NewServerWithClosureAbnormal creates a WS test server with abnormal closure behavior
-func NewServerWithClosureAbnormal(t testing.TB) *Server {
-	return NewServer(t, "/closure-abnormal", getWebsocketHandlerAbnormalClosure())
-}
-
-// NewServerWithEcho creates a WS test server with an echo handler
-func NewServerWithEcho(t testing.TB) *Server {
-	return NewServer(t, "/echo", getWebsocketHandlerEcho())
-}
-
 // NewServer returns a fully configured and running WS test server
-func NewServer(t testing.TB, path string, handler http.Handler) *Server {
+func NewServer(t testing.TB, opts ...func(*Server)) *Server {
 	t.Helper()
 
 	// Create a http.ServeMux and set the httpbin handler as the default
 	mux := http.NewServeMux()
-	mux.Handle(path, handler)
 	mux.Handle("/", httpbin.New().Handler())
 
 	// Initialize the HTTP server and get its details
@@ -137,13 +119,128 @@ func NewServer(t testing.TB, path string, handler http.Handler) *Server {
 		server.Close()
 		cancel()
 	})
-	return &Server{
+	s := &Server{
+		t:             t,
 		Mux:           mux,
 		ServerHTTP:    server,
 		Dialer:        dialer,
 		HTTPTransport: transport,
 		Context:       ctx,
 	}
+	for _, opt := range opts {
+		opt(s)
+	}
+	return s
+}
+
+// WithClosureAbnormalHandler attaches an abnormal closure behavior to Server.
+func WithClosureAbnormalHandler(path string) func(*Server) {
+	handler := func(w http.ResponseWriter, req *http.Request) {
+		conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
+		if err != nil {
+			// TODO: log
+			return
+		}
+		err = conn.Close() // This forces a connection closure without a proper WS close message exchange
+		if err != nil {
+			// TODO: log
+			return
+		}
+	}
+	return func(s *Server) {
+		s.Mux.Handle(path, http.HandlerFunc(handler))
+	}
+}
+
+// NewServerWithEcho attaches an echo handler to Server.
+func WithEchoHandler(path string) func(*Server) {
+	handler := func(w http.ResponseWriter, req *http.Request) {
+		conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
+		if err != nil {
+			return
+		}
+		messageType, r, e := conn.NextReader()
+		if e != nil {
+			return
+		}
+		var wc io.WriteCloser
+		wc, err = conn.NextWriter(messageType)
+		if err != nil {
+			return
+		}
+		if _, err = io.Copy(wc, r); err != nil {
+			return
+		}
+		if err = wc.Close(); err != nil {
+			return
+		}
+		err = conn.WriteControl(websocket.CloseMessage,
+			websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
+			time.Now().Add(10*time.Second),
+		)
+		if err != nil {
+			return
+		}
+	}
+	return func(s *Server) {
+		s.Mux.Handle(path, http.HandlerFunc(handler))
+	}
+}
+
+// WithCDPHandler attaches a a custom CDP handler function to Server.
+func WithCDPHandler(
+	path string,
+	fn func(conn *websocket.Conn, msg *cdproto.Message, writeCh chan cdproto.Message, done chan struct{}),
+	cmdsReceived *[]cdproto.MethodType,
+) func(*Server) {
+	handler := func(w http.ResponseWriter, req *http.Request) {
+		conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
+		if err != nil {
+			return
+		}
+
+		done := make(chan struct{})
+		writeCh := make(chan cdproto.Message)
+
+		// Read
+		go func() {
+			for {
+				select {
+				case <-done:
+					return
+				default:
+				}
+
+				msg, err := CDPReadMsg(conn)
+				if err != nil {
+					close(done)
+					return
+				}
+
+				if msg.Method != "" && cmdsReceived != nil {
+					*cmdsReceived = append(*cmdsReceived, msg.Method)
+				}
+
+				fn(conn, msg, writeCh, done)
+			}
+		}()
+		// Write
+		go func() {
+			for {
+				select {
+				case msg := <-writeCh:
+					CDPWriteMsg(conn, &msg)
+				case <-done:
+					return
+				}
+			}
+		}()
+
+		<-done // Wait for done channel to be closed before closing connection
+	}
+	return func(s *Server) {
+		s.Mux.Handle(path, http.HandlerFunc(handler))
+	}
 }
 
 // TODO: make a websocket.Conn wrapper for CDPxxx methods
@@ -216,98 +313,3 @@ func CDPWriteMsg(conn *websocket.Conn, msg *cdproto.Message) {
 		return
 	}
 }
-
-func getWebsocketHandlerCDP(
-	fn func(conn *websocket.Conn, msg *cdproto.Message, writeCh chan cdproto.Message, done chan struct{}),
-	cmdsReceived *[]cdproto.MethodType) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
-		conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
-		if err != nil {
-			return
-		}
-
-		done := make(chan struct{})
-		writeCh := make(chan cdproto.Message)
-
-		// Read loop
-		go func() {
-			for {
-				select {
-				case <-done:
-					return
-				default:
-				}
-
-				msg, err := CDPReadMsg(conn)
-				if err != nil {
-					close(done)
-					return
-				}
-
-				if msg.Method != "" && cmdsReceived != nil {
-					*cmdsReceived = append(*cmdsReceived, msg.Method)
-				}
-
-				fn(conn, msg, writeCh, done)
-			}
-		}()
-
-		// Write loop
-		go func() {
-			for {
-				select {
-				case msg := <-writeCh:
-					CDPWriteMsg(conn, &msg)
-				case <-done:
-					return
-				}
-			}
-		}()
-
-		<-done // Wait for done channel to be closed before closing connection
-	})
-}
-
-func getWebsocketHandlerEcho() http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
-		conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
-		if err != nil {
-			return
-		}
-		messageType, r, e := conn.NextReader()
-		if e != nil {
-			return
-		}
-		var wc io.WriteCloser
-		wc, err = conn.NextWriter(messageType)
-		if err != nil {
-			return
-		}
-		if _, err = io.Copy(wc, r); err != nil {
-			return
-		}
-		if err = wc.Close(); err != nil {
-			return
-		}
-		err = conn.WriteControl(websocket.CloseMessage,
-			websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
-			time.Now().Add(10*time.Second),
-		)
-		if err != nil {
-			return
-		}
-	})
-}
-
-func getWebsocketHandlerAbnormalClosure() http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
-		conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
-		if err != nil {
-			return
-		}
-		err = conn.Close() // This forces a connection closure without a proper WS close message exchange
-		if err != nil {
-			return
-		}
-	})
-}