Skip to content
This repository was archived by the owner on Jan 30, 2025. It is now read-only.

Commit 6a12e2a

Browse files
inancgumusinancgumus
committed
Refactor WS test server to optional handlers
1 parent ba3d347 commit 6a12e2a

File tree

3 files changed

+136
-134
lines changed

3 files changed

+136
-134
lines changed

‎common/connection_test.go

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ import (
3838
)
3939

4040
func TestConnection(t *testing.T) {
41-
server := ws.NewServerWithEcho(t)
41+
server := ws.NewServer(t, ws.WithEchoHandler("/echo"))
4242

4343
t.Run("connect", func(t *testing.T) {
4444
ctx := context.Background()
@@ -52,7 +52,7 @@ func TestConnection(t *testing.T) {
5252
}
5353

5454
func TestConnectionClosureAbnormal(t *testing.T) {
55-
server := ws.NewServerWithClosureAbnormal(t)
55+
server := ws.NewServer(t, ws.WithClosureAbnormalHandler("/closure-abnormal"))
5656

5757
t.Run("closure abnormal", func(t *testing.T) {
5858
ctx := context.Background()
@@ -69,7 +69,7 @@ func TestConnectionClosureAbnormal(t *testing.T) {
6969
}
7070

7171
func TestConnectionSendRecv(t *testing.T) {
72-
server := ws.NewServerWithCDPHandler(t, ws.CDPDefaultHandler, nil)
72+
server := ws.NewServer(t, ws.WithCDPHandler("/cdp", ws.CDPDefaultHandler, nil))
7373

7474
t.Run("send command with empty reply", func(t *testing.T) {
7575
ctx := context.Background()
@@ -108,19 +108,19 @@ func TestConnectionCreateSession(t *testing.T) {
108108
writeCh <- cdproto.Message{
109109
Method: cdproto.EventTargetAttachedToTarget,
110110
Params: easyjson.RawMessage([]byte(`
111-
{
112-
"sessionId": "0123456789",
113-
"targetInfo": {
114-
"targetId": "abcdef0123456789",
115-
"type": "page",
116-
"title": "",
117-
"url": "about:blank",
118-
"attached": true,
119-
"browserContextId": "0123456789876543210"
120-
},
121-
"waitingForDebugger": false
122-
}
123-
`)),
111+
{
112+
"sessionId": "0123456789",
113+
"targetInfo": {
114+
"targetId": "abcdef0123456789",
115+
"type": "page",
116+
"title": "",
117+
"url": "about:blank",
118+
"attached": true,
119+
"browserContextId": "0123456789876543210"
120+
},
121+
"waitingForDebugger": false
122+
}
123+
`)),
124124
}
125125
writeCh <- cdproto.Message{
126126
ID: msg.ID,
@@ -132,7 +132,7 @@ func TestConnectionCreateSession(t *testing.T) {
132132
}
133133
}
134134

135-
server := ws.NewServerWithCDPHandler(t, handler, &cmdsReceived)
135+
server := ws.NewServer(t, ws.WithCDPHandler("/cdp", handler, &cmdsReceived))
136136

137137
t.Run("create session for target", func(t *testing.T) {
138138
ctx := context.Background()

‎common/session_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ func TestSessionCreateSession(t *testing.T) {
7272
}
7373
}
7474

75-
server := ws.NewServerWithCDPHandler(t, handler, &cmdsReceived)
75+
server := ws.NewServer(t, ws.WithCDPHandler("/cdp", handler, &cmdsReceived))
7676

7777
t.Run("send and recv session commands", func(t *testing.T) {
7878
ctx := context.Background()

‎tests/ws/server.go

Lines changed: 118 additions & 116 deletions
Original file line numberDiff line numberDiff line change
@@ -71,40 +71,22 @@ var (
7171
TargetAttachedToTargetResult = fmt.Sprintf(`{"sessionId":"%s"}`, DummyCDPSessionID)
7272
)
7373

74-
// NewServerWithCDPHandler creates a WS test server with a custom CDP handler function
75-
func NewServerWithCDPHandler(
76-
t testing.TB,
77-
fn func(conn *websocket.Conn, msg *cdproto.Message, writeCh chan cdproto.Message, done chan struct{}),
78-
cmdsReceived *[]cdproto.MethodType) *Server {
79-
return NewServer(t, "/cdp", getWebsocketHandlerCDP(fn, cmdsReceived))
80-
}
81-
8274
// Server can be used as a test alternative to a real CDP compatible browser.
8375
type Server struct {
76+
t testing.TB
8477
Mux *http.ServeMux
8578
ServerHTTP *httptest.Server
8679
Dialer *k6netext.Dialer
8780
HTTPTransport *http.Transport
8881
Context context.Context
8982
}
9083

91-
// NewServerWithClosureAbnormal creates a WS test server with abnormal closure behavior
92-
func NewServerWithClosureAbnormal(t testing.TB) *Server {
93-
return NewServer(t, "/closure-abnormal", getWebsocketHandlerAbnormalClosure())
94-
}
95-
96-
// NewServerWithEcho creates a WS test server with an echo handler
97-
func NewServerWithEcho(t testing.TB) *Server {
98-
return NewServer(t, "/echo", getWebsocketHandlerEcho())
99-
}
100-
10184
// NewServer returns a fully configured and running WS test server
102-
func NewServer(t testing.TB, path string, handler http.Handler) *Server {
85+
func NewServer(t testing.TB, opts ...func(*Server)) *Server {
10386
t.Helper()
10487

10588
// Create a http.ServeMux and set the httpbin handler as the default
10689
mux := http.NewServeMux()
107-
mux.Handle(path, handler)
10890
mux.Handle("/", httpbin.New().Handler())
10991

11092
// Initialize the HTTP server and get its details
@@ -137,13 +119,128 @@ func NewServer(t testing.TB, path string, handler http.Handler) *Server {
137119
server.Close()
138120
cancel()
139121
})
140-
return &Server{
122+
s := &Server{
123+
t: t,
141124
Mux: mux,
142125
ServerHTTP: server,
143126
Dialer: dialer,
144127
HTTPTransport: transport,
145128
Context: ctx,
146129
}
130+
for _, opt := range opts {
131+
opt(s)
132+
}
133+
return s
134+
}
135+
136+
// WithClosureAbnormalHandler attaches an abnormal closure behavior to Server.
137+
func WithClosureAbnormalHandler(path string) func(*Server) {
138+
handler := func(w http.ResponseWriter, req *http.Request) {
139+
conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
140+
if err != nil {
141+
// TODO: log
142+
return
143+
}
144+
err = conn.Close() // This forces a connection closure without a proper WS close message exchange
145+
if err != nil {
146+
// TODO: log
147+
return
148+
}
149+
}
150+
return func(s *Server) {
151+
s.Mux.Handle(path, http.HandlerFunc(handler))
152+
}
153+
}
154+
155+
// NewServerWithEcho attaches an echo handler to Server.
156+
func WithEchoHandler(path string) func(*Server) {
157+
handler := func(w http.ResponseWriter, req *http.Request) {
158+
conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
159+
if err != nil {
160+
return
161+
}
162+
messageType, r, e := conn.NextReader()
163+
if e != nil {
164+
return
165+
}
166+
var wc io.WriteCloser
167+
wc, err = conn.NextWriter(messageType)
168+
if err != nil {
169+
return
170+
}
171+
if _, err = io.Copy(wc, r); err != nil {
172+
return
173+
}
174+
if err = wc.Close(); err != nil {
175+
return
176+
}
177+
err = conn.WriteControl(websocket.CloseMessage,
178+
websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
179+
time.Now().Add(10*time.Second),
180+
)
181+
if err != nil {
182+
return
183+
}
184+
}
185+
return func(s *Server) {
186+
s.Mux.Handle(path, http.HandlerFunc(handler))
187+
}
188+
}
189+
190+
// WithCDPHandler attaches a a custom CDP handler function to Server.
191+
func WithCDPHandler(
192+
path string,
193+
fn func(conn *websocket.Conn, msg *cdproto.Message, writeCh chan cdproto.Message, done chan struct{}),
194+
cmdsReceived *[]cdproto.MethodType,
195+
) func(*Server) {
196+
handler := func(w http.ResponseWriter, req *http.Request) {
197+
conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
198+
if err != nil {
199+
return
200+
}
201+
202+
done := make(chan struct{})
203+
writeCh := make(chan cdproto.Message)
204+
205+
// Read
206+
go func() {
207+
for {
208+
select {
209+
case <-done:
210+
return
211+
default:
212+
}
213+
214+
msg, err := CDPReadMsg(conn)
215+
if err != nil {
216+
close(done)
217+
return
218+
}
219+
220+
if msg.Method != "" && cmdsReceived != nil {
221+
*cmdsReceived = append(*cmdsReceived, msg.Method)
222+
}
223+
224+
fn(conn, msg, writeCh, done)
225+
}
226+
}()
227+
// Write
228+
go func() {
229+
for {
230+
select {
231+
case msg := <-writeCh:
232+
CDPWriteMsg(conn, &msg)
233+
case <-done:
234+
return
235+
}
236+
}
237+
}()
238+
239+
<-done // Wait for done channel to be closed before closing connection
240+
}
241+
return func(s *Server) {
242+
s.Mux.Handle(path, http.HandlerFunc(handler))
243+
}
147244
}
148245

149246
// TODO: make a websocket.Conn wrapper for CDPxxx methods
@@ -216,98 +313,3 @@ func CDPWriteMsg(conn *websocket.Conn, msg *cdproto.Message) {
216313
return
217314
}
218315
}
219-
220-
func getWebsocketHandlerCDP(
221-
fn func(conn *websocket.Conn, msg *cdproto.Message, writeCh chan cdproto.Message, done chan struct{}),
222-
cmdsReceived *[]cdproto.MethodType) http.Handler {
223-
return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
224-
conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
225-
if err != nil {
226-
return
227-
}
228-
229-
done := make(chan struct{})
230-
writeCh := make(chan cdproto.Message)
231-
232-
// Read loop
233-
go func() {
234-
for {
235-
select {
236-
case <-done:
237-
return
238-
default:
239-
}
240-
241-
msg, err := CDPReadMsg(conn)
242-
if err != nil {
243-
close(done)
244-
return
245-
}
246-
247-
if msg.Method != "" && cmdsReceived != nil {
248-
*cmdsReceived = append(*cmdsReceived, msg.Method)
249-
}
250-
251-
fn(conn, msg, writeCh, done)
252-
}
253-
}()
254-
255-
// Write loop
256-
go func() {
257-
for {
258-
select {
259-
case msg := <-writeCh:
260-
CDPWriteMsg(conn, &msg)
261-
case <-done:
262-
return
263-
}
264-
}
265-
}()
266-
267-
<-done // Wait for done channel to be closed before closing connection
268-
})
269-
}
270-
271-
func getWebsocketHandlerEcho() http.Handler {
272-
return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
273-
conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
274-
if err != nil {
275-
return
276-
}
277-
messageType, r, e := conn.NextReader()
278-
if e != nil {
279-
return
280-
}
281-
var wc io.WriteCloser
282-
wc, err = conn.NextWriter(messageType)
283-
if err != nil {
284-
return
285-
}
286-
if _, err = io.Copy(wc, r); err != nil {
287-
return
288-
}
289-
if err = wc.Close(); err != nil {
290-
return
291-
}
292-
err = conn.WriteControl(websocket.CloseMessage,
293-
websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
294-
time.Now().Add(10*time.Second),
295-
)
296-
if err != nil {
297-
return
298-
}
299-
})
300-
}
301-
302-
func getWebsocketHandlerAbnormalClosure() http.Handler {
303-
return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
304-
conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
305-
if err != nil {
306-
return
307-
}
308-
err = conn.Close() // This forces a connection closure without a proper WS close message exchange
309-
if err != nil {
310-
return
311-
}
312-
})
313-
}

0 commit comments

Comments
 (0)