1 package playwright_test
2
3 import (
4 "bytes"
5 "fmt"
6 "log"
7 "net"
8 "net/http"
9 "testing"
10
11 goContext "context"
12
13 "github.com/gorilla/websocket"
14 "github.com/playwright-community/playwright-go"
15 "github.com/stretchr/testify/require"
16 )
17
18 type webSocketServer struct {
19 PORT int
20 server *http.Server
21 }
22
23 func (ws *webSocketServer) Stop() {
24 if err := ws.server.Shutdown(goContext.Background()); err != nil {
25 log.Printf("could not stop ws server: %v", err)
26 }
27 }
28
29 var wsUpgrader = websocket.Upgrader{
30 CheckOrigin: func(r *http.Request) bool { return true },
31 }
32
33 func (ws *webSocketServer) handler(w http.ResponseWriter, r *http.Request) {
34 c, err := wsUpgrader.Upgrade(w, r, nil)
35 if err != nil {
36 log.Print("could not upgrade ws connection:", err)
37 return
38 }
39 if err := c.WriteMessage(websocket.TextMessage, []byte("incoming")); err != nil {
40 log.Println("could not write ws message:", err)
41 return
42 }
43 defer c.Close()
44 for {
45 mt, message, err := c.ReadMessage()
46 if err != nil {
47 if !websocket.IsCloseError(err, websocket.CloseNoStatusReceived) {
48 log.Println("could not read ws message:", err)
49 }
50 break
51 }
52 if bytes.Equal(message, []byte("echo-bin")) {
53 if err := c.WriteMessage(websocket.BinaryMessage, []byte{4, 2}); err != nil {
54 log.Println("could not write ws message:", err)
55 return
56 }
57 break
58 }
59 if bytes.Equal(message, []byte("echo-text")) {
60 if err := c.WriteMessage(mt, []byte("text")); err != nil {
61 log.Println("could not write ws message:", err)
62 return
63 }
64 break
65 }
66 }
67 }
68
69 func newWebsocketServer() *webSocketServer {
70 wsServer := &webSocketServer{
71 PORT: 8012,
72 }
73 mux := http.NewServeMux()
74 mux.HandleFunc("/ws", wsServer.handler)
75 wsServer.server = &http.Server{Handler: mux}
76
77 listener, err := net.Listen("tcp", fmt.Sprintf("localhost:%d", wsServer.PORT))
78 if err != nil {
79 log.Fatal(err)
80 }
81 go func() {
82 if err := wsServer.server.Serve(listener); err != nil && err != http.ErrServerClosed {
83 log.Printf("could not start ws server: %v", err)
84 }
85 }()
86 return wsServer
87 }
88
89 func TestWebSocketShouldWork(t *testing.T) {
90 BeforeEach(t)
91 defer AfterEach(t)
92 wsServer := newWebsocketServer()
93 defer wsServer.Stop()
94 value, err := page.Evaluate(`port => {
95 let cb;
96 const result = new Promise(f => cb = f);
97 const ws = new WebSocket('ws://localhost:' + port + '/ws');
98 ws.addEventListener('message', data => { ws.close(); cb(data.data); });
99 return result;
100 }`, wsServer.PORT)
101 require.NoError(t, err)
102 require.Equal(t, "incoming", value)
103 }
104
105 func TestWebSocketShouldEmitCloseEvents(t *testing.T) {
106 BeforeEach(t)
107 defer AfterEach(t)
108 wsServer := newWebsocketServer()
109 defer wsServer.Stop()
110 wsEvent, err := page.ExpectEvent("websocket", func() error {
111 _, err := page.Evaluate(`port => {
112 const ws = new WebSocket('ws://localhost:' + port + '/ws');
113 ws.addEventListener('message', data => { ws.close() });
114 }`, wsServer.PORT)
115 return err
116 })
117 require.NoError(t, err)
118 ws := wsEvent.(playwright.WebSocket)
119 require.Equal(t, ws.URL(), fmt.Sprintf("ws://localhost:%d/ws", wsServer.PORT))
120 if !ws.IsClosed() {
121 ws.WaitForEvent("close")
122 }
123 require.True(t, ws.IsClosed())
124 }
125
126 func TestWebSocketShouldEmitFrameEvents(t *testing.T) {
127 BeforeEach(t)
128 defer AfterEach(t)
129 wsServer := newWebsocketServer()
130 defer wsServer.Stop()
131
132 sent := [][]byte{}
133 received := [][]byte{}
134
135 page.Once("websocket", func(ws playwright.WebSocket) {
136 ws.On("framesent", func(payload []byte) {
137 sent = append(sent, payload)
138 })
139 ws.On("framereceived", func(payload []byte) {
140 received = append(received, payload)
141 })
142 })
143 wsEvent, err := page.ExpectEvent("websocket", func() error {
144 _, err := page.Evaluate(`port => {
145 const ws = new WebSocket('ws://localhost:' + port + '/ws');
146 ws.addEventListener('open', () => {
147 ws.send('echo-text');
148 });
149 }`, wsServer.PORT)
150 return err
151 })
152 require.NoError(t, err)
153 ws := wsEvent.(playwright.WebSocket)
154 if !ws.IsClosed() {
155 ws.WaitForEvent("close")
156 }
157
158 require.Equal(t, sent, [][]byte{[]byte("echo-text")})
159 require.Equal(t, received, [][]byte{[]byte("incoming"), []byte("text")})
160 }
161
162 func TestWebSocketShouldEmitBinaryFrameEvents(t *testing.T) {
163 BeforeEach(t)
164 defer AfterEach(t)
165 wsServer := newWebsocketServer()
166 defer wsServer.Stop()
167
168 sent := [][]byte{}
169 received := [][]byte{}
170
171 page.Once("websocket", func(ws playwright.WebSocket) {
172 ws.On("framesent", func(payload []byte) {
173 sent = append(sent, payload)
174 })
175 ws.On("framereceived", func(payload []byte) {
176 received = append(received, payload)
177 })
178 })
179 wsEvent, err := page.ExpectEvent("websocket", func() error {
180 _, err := page.Evaluate(`port => {
181 const ws = new WebSocket('ws://localhost:' + port + '/ws');
182 ws.addEventListener('open', () => {
183 const binary = new Uint8Array(5);
184 for (let i = 0; i < 5; ++i)
185 binary[i] = i;
186 ws.send(binary);
187 ws.send('echo-bin');
188 });
189 }`, wsServer.PORT)
190 return err
191 })
192 require.NoError(t, err)
193 ws := wsEvent.(playwright.WebSocket)
194 if !ws.IsClosed() {
195 ws.WaitForEvent("close")
196 }
197
198 require.Equal(t, sent, [][]byte{{0, 1, 2, 3, 4}, []byte("echo-bin")})
199 require.Equal(t, received, [][]byte{[]byte("incoming"), {4, 2}})
200 }
201
View as plain text