test: unflake test_should_emit_frame_events test (microsoft#781)

mxschmitt · web-flow · commit 9ec64c6f7145 · 2021-06-28T21:22:28.000+02:00
diff --git a/playwright/_impl/_network.py b/playwright/_impl/_network.py
@@ -341,13 +341,13 @@ async def wait_for_event(
     def _on_frame_sent(self, opcode: int, data: str) -> None:
         if opcode == 2:
             self.emit(WebSocket.Events.FrameSent, base64.b64decode(data))
-        else:
+        elif opcode == 1:
             self.emit(WebSocket.Events.FrameSent, data)
 
     def _on_frame_received(self, opcode: int, data: str) -> None:
         if opcode == 2:
             self.emit(WebSocket.Events.FrameReceived, base64.b64decode(data))
-        else:
+        elif opcode == 1:
             self.emit(WebSocket.Events.FrameReceived, data)
 
     def is_closed(self) -> bool:
diff --git a/tests/async/test_websocket.py b/tests/async/test_websocket.py
@@ -55,30 +55,32 @@ async def test_should_emit_close_events(page, ws_server):
 
 
 async def test_should_emit_frame_events(page, ws_server):
-    sent = []
-    received = []
+    log = []
+    socke_close_future = asyncio.Future()
 
     def on_web_socket(ws):
-        ws.on("framesent", lambda payload: sent.append(payload))
-        ws.on("framereceived", lambda payload: received.append(payload))
+        log.append("open")
+        ws.on("framesent", lambda payload: log.append(f"sent<{payload}>"))
+        ws.on("framereceived", lambda payload: log.append(f"received<{payload}>"))
+        ws.on(
+            "close", lambda: (log.append("close"), socke_close_future.set_result(None))
+        )
 
     page.on("websocket", on_web_socket)
-    async with page.expect_event("websocket") as ws_info:
+    async with page.expect_event("websocket"):
         await page.evaluate(
             """port => {
             const ws = new WebSocket('ws://localhost:' + port + '/ws');
-            ws.addEventListener('open', () => {
-                ws.send('echo-text');
-            });
+            ws.addEventListener('open', () => ws.send('outgoing'));
+            ws.addEventListener('message', () => ws.close())
         }""",
             ws_server.PORT,
         )
-    ws = await ws_info.value
-    if not ws.is_closed():
-        await ws.wait_for_event("close")
-
-    assert sent == ["echo-text"]
-    assert received == ["incoming", "text"]
+    await socke_close_future
+    assert log[0] == "open"
+    assert log[3] == "close"
+    log.sort()
+    assert log == ["close", "open", "received<incoming>", "sent<outgoing>"]
 
 
 async def test_should_emit_binary_frame_events(page, ws_server):
