neon-testing 1.0.2 → 1.1.1

package/README.md

@@ -1,15 +1,18 @@
 # Neon testing
 
-A Vitest utility for automated integration tests with [Neon](https://neon.com/).
+[![Integration tests](https://github.com/starmode-base/neon-testing/actions/workflows/test.yml/badge.svg)](https://github.com/starmode-base/neon-testing/actions/workflows/test.yml)
 
-Each test file runs against its own isolated PostgreSQL database (Neon branch), ensuring clean, parallel, and reproducible testing of code that relies on a database. Because it uses a real database, you can test code logic that depends on database features such as transaction rollbacks, unique constraints, and more.
+A [Vitest](https://vitest.dev/) utility for seamless integration tests with [Neon Postgres](https://neon.com/).
 
-Using an actual clone of your production database for integration testing lets you verify functionality that mocks cannot.
+Each test file runs against its own isolated PostgreSQL database (Neon branch), ensuring clean, parallel, and reproducible testing of code that interacts with a database. Because it uses a real, isolated clone of your production database, you can test code logic that depends on database features, such as transaction rollbacks, unique constraints, and more.
+
+**Testing against a clone of your production database lets you verify functionality that mocks cannot.**
 
 ## Features
 
 - 🔄 **Isolated test environments** - Each test file runs against its own Postgres database with your actual schema and constraints
 - 🧹 **Automatic cleanup** - Neon test branches are created and destroyed automatically
+- 🐛 **Debug friendly** - Option to preserve test branches for debugging failed tests
 - 🛡️ **TypeScript native** - No JavaScript support
 - 🎯 **ESM only** - No CommonJS support
 
@@ -20,12 +23,20 @@ Using an actual clone of your production database for integration testing lets y
 1. **Test execution**: Your tests run against the isolated database
 1. **Cleanup**: After tests complete, the branch is automatically deleted
 
+### Test isolation
+
+Tests within a test file share the same database instance (Neon branch), so while all test files are isolated, tests within a test file are intentionally not.
+
+This works because Vitest runs test files in parallel, but tests within each test file run sequentially one at a time.
+
+If you prefer individual tests within a test file to be isolated, [simply clean up the database in a beforeEach lifecycle](examples/isolated.test.ts).
+
 ## Quick start
 
 ### Prerequisites
 
 - A [Neon project](https://console.neon.tech/app/projects) with a database
-- A [Neon API key](https://neon.tech/docs/manage/api-keys) for programmatic access
+- A [Neon API key](https://neon.com/docs/manage/api-keys) for programmatic access
 
 ### Install
 
@@ -36,13 +47,18 @@ bun add -d neon-testing
 ### Minimal example
 
 ```typescript
-// database.test.ts
+// minimal.test.ts
 import { expect, test } from "vitest";
 import { makeNeonTesting } from "neon-testing";
 import { Pool } from "@neondatabase/serverless";
 
 // Enable Neon test branch for this test file
-makeNeonTesting({ apiKey: "apiKey", projectId: "projectId" })();
+makeNeonTesting({
+  apiKey: "apiKey",
+  projectId: "projectId",
+  // Recommended for Neon WebSocket drivers to automatically close connections
+  autoCloseWebSockets: true,
+})();
 
 test("database operations", async () => {
   const pool = new Pool({ connectionString: process.env.DATABASE_URL });
@@ -52,8 +68,6 @@ test("database operations", async () => {
 
   const users = await pool.query(`SELECT * FROM users`);
   expect(users.rows).toStrictEqual([{ id: 1, name: "Ellen Ripley" }]);
-
-  await pool.end();
 });
 ```
 
@@ -74,20 +88,21 @@ export const withNeonTestBranch = makeNeonTesting({
 });
 ```
 
-See all available options in [NeonTestingOptions](https://github.com/starmode-base/neon-testing/blob/main/index.ts#L30-L41).
-
 #### 2. Enable database testing
 
 Then call the exported test lifecycle function in the test files where you need database access.
 
 ```typescript
-// database.test.ts
+// recommended.test.ts
 import { expect, test } from "vitest";
 import { withNeonTestBranch } from "./test-setup";
 import { Pool } from "@neondatabase/serverless";
 
 // Enable Neon test branch for this test file
-withNeonTestBranch();
+withNeonTestBranch({
+  // Recommended for Neon WebSocket drivers to automatically close connections
+  autoCloseWebSockets: true,
+});
 
 test("database operations", async () => {
   const pool = new Pool({ connectionString: process.env.DATABASE_URL });
@@ -97,14 +112,51 @@ test("database operations", async () => {
 
   const users = await pool.query(`SELECT * FROM users`);
   expect(users.rows).toStrictEqual([{ id: 1, name: "Ellen Ripley" }]);
+});
+```
+
+## Drivers
+
+This library works with any database driver that supports Neon Postgres and Vitest. The examples below demonstrate connection management, transaction support, and test isolation patterns for some popular drivers.
+
+**IMPORTANT:** For [Neon WebSocket drivers](https://neon.com/docs/serverless/serverless-driver), enable `autoCloseWebSockets` in your `makeNeonTesting()` or `withNeonTestBranch()` configuration. This automatically closes WebSocket connections when deleting test branches, preventing connection termination errors.
+
+### Examples
+
+- [Neon serverless WebSocket](examples/drivers/ws-neon.test.ts)
+- [Neon serverless WebSocket + Drizzle](examples/drivers/ws-neon-drizzle.test.ts)
+- [Neon serverless HTTP](examples/drivers/http-neon.test.ts)
+- [Neon serverless HTTP + Drizzle](examples/drivers/http-neon-drizzle.test.ts)
+- [node-postgres](examples/drivers/tcp-pg.test.ts)
+- [node-postgres + Drizzle](examples/drivers/tcp-pg-drizzle.test.ts)
+- [Postgres.js](examples/drivers/tcp-postgres.test.ts)
+- [Postgres.js + Drizzle](examples/drivers/tcp-postgres-drizzle.test.ts)
+
+## Configuration
+
+You configure neon-testing in two places:
+
+- **Base settings** in `makeNeonTesting()`
+- **Optional overrides** in `withNeonTestBranch()`
+
+See all available options in [NeonTestingOptions](https://github.com/starmode-base/neon-testing/blob/main/index.ts#L33-L75).
+
+### Base configuration
+
+Configure the base settings in `makeNeonTesting()`:
+
+```typescript
+import { makeNeonTesting } from "neon-testing";
 
-  await pool.end();
+export const withNeonTestBranch = makeNeonTesting({
+  apiKey: "apiKey",
+  projectId: "projectId",
 });
 ```
 
-#### Override configuration
+### Override configuration
 
-Branch from a specific branch instead of main:
+Override the base configuration in specific test files with `withNeonTestBranch()`:
 
 ```typescript
 import { withNeonTestBranch } from "./test-setup";
@@ -112,34 +164,67 @@ import { withNeonTestBranch } from "./test-setup";
 withNeonTestBranch({ parentBranchId: "br-staging-123" });
 ```
 
-Don't copy data when branching:
+## Utilities
+
+### deleteAllTestBranches()
+
+The `deleteAllTestBranches()` function is a utility that deletes all test branches from your Neon project. This is useful for cleanup when tests fail unexpectedly and leave orphaned test branches.
 
 ```typescript
 import { withNeonTestBranch } from "./test-setup";
 
-withNeonTestBranch({ schemaOnly: true });
+// Access the cleanup utility
+await withNeonTestBranch.deleteAllTestBranches();
 ```
 
-See all available options in [NeonTestingOptions](https://github.com/starmode-base/neon-testing/blob/main/index.ts#L30-L41).
+The function identifies test branches by looking for the `integration-test: true` annotation that neon-testing automatically adds to all test branches it creates.
+
+### lazySingleton()
+
+The `lazySingleton()` function creates a lazy singleton from a factory function. This is useful for managing database connections efficiently:
 
-## Isolate individual tests
+```typescript
+import { lazySingleton } from "neon-testing";
+import { neon } from "@neondatabase/serverless";
 
-Tests within a single test file share the same database instance (Neon branch), so while all test files are isolated, tests within a test file are not. If you prefer individual tests within a test file to be isolated, [simply clean up the database in a beforeEach lifecycle](examples/neon-serverless-http-isolated.test.ts).
+const sql = lazySingleton(() => neon(process.env.DATABASE_URL!));
 
-This works because Vitest runs test files in parallel, but tests within each test file run one at a time.
+// The connection is only created when first called
+test("database operations", async () => {
+  const users = await sql()`SELECT * FROM users`;
+  // ...
+});
+```
 
 ## Contributing
 
 Contributions are welcome! Please open issues or pull requests on [GitHub](https://github.com/starmode-base/neon-testing/pulls).
 
+### Environment
+
+To run tests locally, create an `.env` file in the project root with these keys:
+
+- `NEON_API_KEY="***"`
+- `NEON_PROJECT_ID="***"`
+
+Create a free Neon project at [neon.com](https://neon.com/) to test with.
+
+### Release
+
+To make a new release, run:
+
+```sh
+bun run release
+```
+
+The command will abort if there are uncommitted changes in the working tree, or if the `version` in [package.json](package.json) has not been incremented.
+
 ## License
 
 This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
 
-## Need expert help?
-
-Hi, I'm [@lirbank](https://github.com/lirbank). I take on a few consulting projects each year where I help companies build, unblock, and ship. Here's what I do:
+## Need help?
 
-**[STΛR MODΞ](https://www.starmode.dev/)** — A boutique AI development studio I run with AI/ML expert and data scientist [@spencer-g-smith](https://github.com/spencer-g-smith). We help companies build accurate AI solutions: AI-first apps, advanced workflows, and agentic systems.
+Hi, I’m [Mikael Lirbank](https://www.lirbank.com/). I help teams build reliable AI systems. I care about quality—AI evals, robust test suites, clean data models, and clean architecture. Sometimes I draw user interfaces.
 
-**[Mikael Lirbank](https://www.lirbank.com/)** — My solo practice, focused on web app development, test automation, code quality, and technical architecture. I'm friendly and happy to help with the hard stuff.
+Want to ship faster without breaking things? Let’s talk.

package/index.ts

@@ -7,6 +7,9 @@ import {
   type ConnectionDetails,
 } from "@neondatabase/api-client";
 import { afterAll, beforeAll } from "vitest";
+import { neonConfig } from "@neondatabase/serverless";
+
+export { lazySingleton } from "./singleton";
 
 /**
  * Creates a PostgreSQL connection URI from connection parameters
@@ -28,16 +31,47 @@ function createConnectionUri(
 }
 
 export interface NeonTestingOptions {
-  /** The Neon API key, this is used to create and teardown test branches */
+  /**
+   * The Neon API key, this is used to create and teardown test branches
+   *
+   * https://neon.com/docs/manage/api-keys#creating-api-keys
+   */
   apiKey: string;
-  /** The Neon project ID to operate on */
+  /**
+   * The Neon project ID to operate on
+   *
+   * https://console.neon.tech/app/projects
+   */
   projectId: string;
-  /** The parent branch ID for the new branch */
+  /**
+   * The parent branch ID for the new branch. If omitted or empty, the branch
+   * will be created from the project's default branch.
+   */
   parentBranchId?: string;
-  /** Whether to create a schema-only branch (default: false) */
+  /**
+   * Whether to create a schema-only branch (default: false)
+   */
   schemaOnly?: boolean;
-  /** The type of connection to create (pooler is recommended) */
+  /**
+   * The type of connection to create (pooler is recommended)
+   */
   endpoint?: "pooler" | "direct";
+  /**
+   * Delete the test branch in afterAll (default: true)
+   *
+   * Disabling this will leave each test branch in the Neon project after the
+   * test suite runs
+   */
+  deleteBranch?: boolean;
+  /**
+   * Automatically close Neon WebSocket connections opened during tests before
+   * deleting the branch (default: false)
+   *
+   * Suppresses the specific Neon WebSocket "Connection terminated unexpectedly"
+   * error that may surface when deleting a branch with open WebSocket
+   * connections
+   */
+  autoCloseWebSockets?: boolean;
 }
 
 /** Options for overriding test database setup (excludes apiKey) */
@@ -47,11 +81,12 @@ export type NeonTestingOverrides = Omit<Partial<NeonTestingOptions>, "apiKey">;
  * Factory function that creates a Neon test database setup/teardown function
  * for Vitest test suites.
  *
- * @param apiKey - The Neon API key
- * @param projectId - The Neon project ID
- * @param endpoint - The type of connection to create (pooler or direct)
- * @param parentBranchId - The parent branch ID for the new branch
- * @param schemaOnly - Whether to create a schema-only branch
+ * @param apiKey - The Neon API key, this is used to create and teardown test branches
+ * @param projectId - The Neon project ID to operate on
+ * @param parentBranchId - The parent branch ID for the new branch. If omitted or empty, the branch will be created from the project's default branch.
+ * @param schemaOnly - Whether to create a schema-only branch (default: false)
+ * @param endpoint - The type of connection to create (pooler is recommended)
+ * @param deleteBranch - Delete the test branch in afterAll (default: true). Disabling this will leave each test branch in the Neon project after the test suite runs
  * @returns A setup/teardown function for Vitest test suites
  *
  * Side effects:
@@ -93,6 +128,21 @@ export function makeNeonTesting(factoryOptions: NeonTestingOptions) {
     // Each test file gets its own branch ID and database client
     let branchId: string | undefined;
 
+    // List of tracked Neon WebSocket connections
+    const neonSockets = new Set<WebSocket>();
+
+    // Custom WebSocket constructor that tracks Neon WebSocket connections
+    class TrackingWebSocket extends WebSocket {
+      constructor(url: string) {
+        super(url);
+
+        // Only track Neon WebSocket connections
+        if (!url.includes(".neon.tech/")) return;
+
+        neonSockets.add(this);
+      }
+    }
+
     /**
      * Create a new test branch
      *
@@ -114,6 +164,7 @@ export function makeNeonTesting(factoryOptions: NeonTestingOptions) {
       branchId = data.branch.id;
 
       const [connectionUri] = data.connection_uris ?? [];
+
       if (!connectionUri) {
         throw new Error("No connection URI found");
       }
@@ -134,12 +185,33 @@ export function makeNeonTesting(factoryOptions: NeonTestingOptions) {
     }
 
     beforeAll(async () => {
-      process.env.DATABASE_URL = await createBranch();
+      process.env.DATABASE_URL = await withRetry(createBranch, {
+        maxRetries: 5,
+        baseDelayMs: 1000,
+      });
+
+      if (options.autoCloseWebSockets) {
+        // Install a custom WebSocket constructor that tracks Neon WebSocket
+        // connections and closes them before deleting the branch
+        neonConfig.webSocketConstructor = TrackingWebSocket;
+      }
     });
 
     afterAll(async () => {
-      await deleteBranch();
       process.env.DATABASE_URL = undefined;
+
+      // Close all tracked Neon WebSocket connections before deleting the branch
+      if (options.autoCloseWebSockets) {
+        // Suppress Neon WebSocket "Connection terminated unexpectedly" error
+        process.prependListener("uncaughtException", neonWsErrorHandler);
+
+        // Close tracked Neon WebSocket connections before deleting the branch
+        neonSockets.forEach((ws) => ws.close());
+      }
+
+      if (options.deleteBranch !== false) {
+        await deleteBranch();
+      }
     });
   };
 
@@ -148,3 +220,66 @@ export function makeNeonTesting(factoryOptions: NeonTestingOptions) {
 
   return testDbSetup;
 }
+
+/**
+ * Error handler: Suppress Neon WebSocket "Connection terminated unexpectedly"
+ * error
+ */
+const neonWsErrorHandler = (error: Error) => {
+  const isNeonWsClose =
+    error.message.includes("Connection terminated unexpectedly") &&
+    error.stack?.includes("@neondatabase/serverless");
+
+  if (isNeonWsClose) {
+    // Swallow this specific Neon WS termination error
+    return;
+  }
+
+  // For any other error, detach and rethrow
+  throw error;
+};
+
+/**
+ * Reusable API call wrapper with automatic retry on 423 errors with exponential
+ * backoff
+ *
+ * https://neon.com/docs/reference/typescript-sdk#error-handling
+ * https://neon.com/docs/changelog/2022-07-20
+ */
+async function withRetry<T>(
+  fn: () => Promise<T>,
+  options: {
+    maxRetries: number;
+    baseDelayMs: number;
+  },
+): Promise<T> {
+  if (!Number.isInteger(options.maxRetries) || options.maxRetries <= 0) {
+    throw new Error("maxRetries must be a positive integer");
+  }
+
+  if (!Number.isInteger(options.baseDelayMs) || options.baseDelayMs <= 0) {
+    throw new Error("baseDelayMs must be a positive integer");
+  }
+
+  for (let attempt = 1; attempt <= options.maxRetries; attempt++) {
+    try {
+      return await fn();
+    } catch (error: any) {
+      const status = error?.response?.status;
+
+      if (status === 423 && attempt < options.maxRetries) {
+        const delay = options.baseDelayMs * Math.pow(2, attempt - 1);
+
+        console.log(
+          `API call failed with 423, retrying in ${delay}ms (attempt ${attempt}/${options.maxRetries})`,
+        );
+        await new Promise((resolve) => setTimeout(resolve, delay));
+
+        continue;
+      }
+
+      throw error;
+    }
+  }
+  throw new Error("apiCallWithRetry reached unexpected end");
+}

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "neon-testing",
-  "version": "1.0.2",
-  "description": "A Vitest utility for automated integration tests with Neon",
+  "version": "1.1.1",
+  "description": "A Vitest utility for seamless integration tests with Neon Postgres",
   "keywords": [
     "neon",
     "postgres",
@@ -17,7 +17,8 @@
   "module": "index.ts",
   "type": "module",
   "files": [
-    "index.ts"
+    "index.ts",
+    "singleton.ts"
   ],
   "scripts": {
     "test": "vitest",
@@ -27,17 +28,19 @@
     "postpublish": "git tag v$(bun -p \"require('./package.json').version\") && git push --tags"
   },
   "dependencies": {
-    "@neondatabase/api-client": "^2.1.0"
+    "@neondatabase/api-client": "^2.2.0"
   },
   "peerDependencies": {
-    "typescript": "^5.8.3",
     "vitest": "^3.2.4"
   },
   "devDependencies": {
     "@neondatabase/serverless": "^1.0.1",
-    "dotenv": "^16.5.0",
-    "pg": "^8.16.2",
+    "dotenv": "^17.2.1",
+    "drizzle-orm": "^0.44.4",
+    "pg": "^8.16.3",
     "postgres": "^3.4.7",
-    "prettier": "^3.5.3"
+    "prettier": "^3.6.2",
+    "typescript": "^5.9.2",
+    "vitest": "^3.2.4"
   }
 }

package/singleton.ts

@@ -0,0 +1,10 @@
+/**
+ * Create a lazy singleton from a factory function
+ */
+export function lazySingleton<T>(factory: () => T): () => T {
+  let instance: T | undefined;
+  return () => {
+    instance ??= factory();
+    return instance;
+  };
+}