From 80d3092b0526df22f42209cd27cb5825754aa901 Mon Sep 17 00:00:00 2001
From: Brendan Burns <5751682+brendandburns@users.noreply.github.com>
Date: Sat, 22 Feb 2025 23:40:34 +0000
Subject: [PATCH] Add coverage improvements for web-socket-handler.

---
 src/web-socket-handler.ts      |  11 +++-
 src/web-socket-handler_test.ts | 115 +++++++++++++++++++++++++++++----
 2 files changed, 113 insertions(+), 13 deletions(-)

diff --git a/src/web-socket-handler.ts b/src/web-socket-handler.ts
index 6772427dc1b..59592011556 100644
--- a/src/web-socket-handler.ts
+++ b/src/web-socket-handler.ts
@@ -137,11 +137,12 @@ export class WebSocketHandler implements WebSocketInterface {
         stdin: stream.Readable,
         streamNum: number = 0,
         retryCount: number = 3,
+        // kind of hacky, but otherwise we can't wait for the writes to flush before testing.
+        addFlushForTesting: boolean = false,
     ): () => WebSocket.WebSocket | null {
         if (retryCount < 0) {
             throw new Error("retryCount can't be lower than 0.");
         }
-
         let queue: Promise<void> = Promise.resolve();
         let ws: WebSocket.WebSocket | null = null;
 
@@ -158,8 +159,14 @@ export class WebSocketHandler implements WebSocketInterface {
             });
         });
 
+        if (addFlushForTesting) {
+            stdin.on('flush', async () => {
+                await queue;
+            });
+        }
+
         stdin.on('end', () => {
-            if (ws) {
+            if (ws !== null) {
                 ws.close();
             }
         });
diff --git a/src/web-socket-handler_test.ts b/src/web-socket-handler_test.ts
index 84f7b6512e9..0a0ae922b60 100644
--- a/src/web-socket-handler_test.ts
+++ b/src/web-socket-handler_test.ts
@@ -1,4 +1,4 @@
-import { deepStrictEqual, notStrictEqual, rejects, strictEqual, throws } from 'node:assert';
+import { deepStrictEqual, equal, notStrictEqual, rejects, strictEqual, throws } from 'node:assert';
 import { Readable, Writable } from 'node:stream';
 import { setImmediate as setImmediatePromise } from 'node:timers/promises';
 import WebSocket from 'isomorphic-ws';
@@ -335,7 +335,9 @@ describe('V5 protocol support', () => {
             protocol: 'v5.channel.k8s.io',
         } as WebSocket.WebSocket;
         let uriOut = '';
-        let endCalled = false;
+        let stderrEndCalled = false;
+        let stdoutEndCalled = false;
+        let stdinPauseCalled = false;
         const handler = new WebSocketHandler(
             kc,
             (uri: string, protocols: string[], opts: WebSocket.ClientOptions): WebSocket.WebSocket => {
@@ -343,11 +345,20 @@ describe('V5 protocol support', () => {
                 return mockWs as WebSocket.WebSocket;
             },
             {
-                stdin: process.stdin,
-                stderr: process.stderr,
+                stdin: {
+                    pause: () => {
+                        stdinPauseCalled = true;
+                        return {} as Readable;
+                    },
+                } as Readable,
+                stderr: {
+                    end: () => {
+                        stderrEndCalled = true;
+                    },
+                } as Writable,
                 stdout: {
                     end: () => {
-                        endCalled = true;
+                        stdoutEndCalled = true;
                     },
                 } as Writable,
             },
@@ -364,17 +375,30 @@ describe('V5 protocol support', () => {
             type: 'open',
         };
         mockWs.onopen!(event);
-        const closeBuff = Buffer.alloc(2);
-        closeBuff.writeUint8(255, 0);
-        closeBuff.writeUint8(WebSocketHandler.StdoutStream, 1);
-
+        // Close stdin/stdout with Buffers
+        [WebSocketHandler.StdinStream, WebSocketHandler.StdoutStream].forEach((stream) => {
+            const closeBuff = Buffer.alloc(2);
+            closeBuff.writeUint8(255, 0);
+            closeBuff.writeUint8(stream, 1);
+
+            mockWs.onmessage!({
+                data: closeBuff,
+                type: 'type',
+                target: mockWs,
+            });
+        });
+        // Close stderr with a string \xff is 'close' \x02 is the stderr stream number
+        // so that both paths are tested.
+        const closeMsg = '\xFF\x02';
         mockWs.onmessage!({
-            data: closeBuff,
+            data: closeMsg,
             type: 'type',
             target: mockWs,
         });
         await promise;
-        strictEqual(endCalled, true);
+        strictEqual(stdoutEndCalled, true);
+        strictEqual(stderrEndCalled, true);
+        strictEqual(stdinPauseCalled, true);
     });
     it('should handle closing stdin < v4 protocol', () => {
         const ws = {
@@ -436,4 +460,73 @@ describe('Restartable Handle Standard Input', () => {
             strictEqual(count, retryTimes);
         });
     });
+
+    it('should work correctly', async () => {
+        let sent: Buffer | null = null;
+        const ws = {
+            protocol: 'v5.channel.k8s.io',
+            send: (data) => {
+                sent = data;
+            },
+            readyState: WebSocket.OPEN,
+            close: () => {
+                throw new Error('should not be called');
+            },
+        } as unknown as WebSocket;
+        const p = new Promise<WebSocket>((resolve, reject) => resolve(ws));
+        let dataCb: any;
+        let endCb: any;
+        let flushCb: any;
+
+        const r = {
+            on: (verb, cb) => {
+                if (verb === 'data') {
+                    dataCb = cb;
+                }
+                if (verb === 'end') {
+                    endCb = cb;
+                }
+                if (verb == 'flush') {
+                    flushCb = cb;
+                }
+            },
+        } as Readable;
+
+        WebSocketHandler.restartableHandleStandardInput(() => p, r, 0, 4, true);
+
+        dataCb('some test data');
+        endCb();
+        await flushCb();
+
+        equal(sent!.toString(), '\x00some test data');
+    });
+
+    it('should work if the web socket exists', () => {
+        let sent: Buffer | null = null;
+        const ws = {
+            protocol: 'v5.channel.k8s.io',
+            send: (data) => {
+                sent = data;
+            },
+            readyState: WebSocket.OPEN,
+            close: () => {
+                throw new Error('should not be called');
+            },
+        } as unknown as WebSocket;
+        let count = 0;
+        WebSocketHandler.processData(
+            'some test data',
+            ws,
+            (): Promise<WebSocket.WebSocket> => {
+                return new Promise<WebSocket.WebSocket>((resolve) => {
+                    count++;
+                    resolve(ws as WebSocket.WebSocket);
+                });
+            },
+            0,
+            5,
+        );
+        equal(sent!.toString(), '\x00some test data');
+        strictEqual(count, 0);
+    });
 });