From 6af67df4574f6d654e3f69fb7dd58369d19a7417 Mon Sep 17 00:00:00 2001
From: George MacKerron <george@mackerron.co.uk>
Date: Wed, 18 Dec 2024 17:28:26 +0000
Subject: [PATCH] Refactired out WebSocket shimming

---
 tests/http.test.ts |  8 ++------
 tests/ws.test.ts   |  8 ++------
 tests/ws.ts        | 11 +++++++++++
 3 files changed, 15 insertions(+), 12 deletions(-)
 create mode 100644 tests/ws.ts

diff --git a/tests/http.test.ts b/tests/http.test.ts
index a79e87a..88a4b86 100644
--- a/tests/http.test.ts
+++ b/tests/http.test.ts
@@ -1,18 +1,14 @@
 import { expect, test, vi, beforeAll } from 'vitest';
 import { neon, neonConfig, Pool } from '../dist/npm';
 import { sampleQueries } from './sampleQueries';
+import { shimWebSocket } from './ws';
 
 const DB_URL = process.env.VITE_NEON_DB_URL!;
 const sql = neon(DB_URL);
 const sqlFull = neon(DB_URL, { fullResults: true });
 const pool = new Pool({ connectionString: DB_URL });
 
-beforeAll(async () => {
-  if (typeof WebSocket !== 'function') {
-    const { WebSocket } = await import('ws');
-    neonConfig.webSocketConstructor = WebSocket;
-  }
-});
+beforeAll(shimWebSocket);
 
 test(
   'http query results match WebSocket query results',
diff --git a/tests/ws.test.ts b/tests/ws.test.ts
index 5b9a0fd..cf2092f 100644
--- a/tests/ws.test.ts
+++ b/tests/ws.test.ts
@@ -2,6 +2,7 @@ import { expect, test, beforeAll } from 'vitest';
 import { Pool as PgPool } from 'pg';
 import * as subtls from 'subtls';
 import { sampleQueries } from './sampleQueries';
+import { shimWebSocket } from './ws';
 import {
   neon,
   neonConfig,
@@ -33,12 +34,7 @@ function functionsToPlaceholders(x: any) {
   );
 }
 
-beforeAll(async () => {
-  if (typeof WebSocket !== 'function') {
-    const { WebSocket } = await import('ws');
-    neonConfig.webSocketConstructor = WebSocket;
-  }
-});
+beforeAll(shimWebSocket);
 
 test(
   'WebSockets query results match pg/TCP query results, using pool.connect()',
diff --git a/tests/ws.ts b/tests/ws.ts
new file mode 100644
index 0000000..99c2159
--- /dev/null
+++ b/tests/ws.ts
@@ -0,0 +1,11 @@
+import { neonConfig } from '../dist/npm';
+
+export async function shimWebSocket() {
+  if (typeof WebSocket !== 'function') {
+    console.info('Native WebSocket not found: importing ws for neonConfig.webSocketConstructor');
+    const { WebSocket } = await import('ws');
+    neonConfig.webSocketConstructor = WebSocket;
+  } else {
+    console.info('Using native WebSocket');
+  }
+}