From 043c5aaffb2248a80946cbcd4563952a51d1a144 Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Fri, 16 Feb 2024 18:16:45 -0800
Subject: [PATCH] chore: follow up to web socket util extraction (#29536)

---
 packages/playwright-core/src/utils/wsServer.ts | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/packages/playwright-core/src/utils/wsServer.ts b/packages/playwright-core/src/utils/wsServer.ts
index e510a7ae8e..63c0ed1fed 100644
--- a/packages/playwright-core/src/utils/wsServer.ts
+++ b/packages/playwright-core/src/utils/wsServer.ts
@@ -39,10 +39,10 @@ export type WSConnection = {
 };
 
 export type WSServerDelegate = {
-  onHeaders: (headers: string[]) => void;
-  onUpgrade: (request: http.IncomingMessage, socket: stream.Duplex) => { error: string } | undefined;
+  onHeaders?: (headers: string[]) => void;
+  onUpgrade?: (request: http.IncomingMessage, socket: stream.Duplex) => { error: string } | undefined;
   onConnection: (request: http.IncomingMessage, url: URL, ws: WebSocket, id: string) => WSConnection;
-  onClose(): Promise<void>;
+  onClose?(): Promise<void>;
 };
 
 export class WSServer {
@@ -89,7 +89,8 @@ export class WSServer {
       perMessageDeflate,
     });
 
-    this._wsServer.on('headers', headers => this._delegate.onHeaders(headers));
+    if (this._delegate.onHeaders)
+      this._wsServer.on('headers', headers => this._delegate.onHeaders!(headers));
 
     server.on('upgrade', (request, socket, head) => {
       const pathname = new URL('http://localhost' + request.url!).pathname;
@@ -98,7 +99,7 @@ export class WSServer {
         socket.destroy();
         return;
       }
-      const upgradeResult = this._delegate.onUpgrade(request, socket);
+      const upgradeResult = this._delegate.onUpgrade?.(request, socket);
       if (upgradeResult) {
         socket.write(upgradeResult.error);
         socket.destroy();
@@ -142,5 +143,7 @@ export class WSServer {
     this._wsServer = undefined;
     this.server = undefined;
     debugLogger.log('server', 'closed server');
+
+    await this._delegate.onClose?.();
   }
 }