Refactired out WebSocket shimming

neondatabase · Dec 18, 2024 · 6af67df · 6af67df
1 parent bcf0060
commit 6af67df
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 12 deletions.
diff --git 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
@@ -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
@@ -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');
+  }
+}