tests: Add tests for client auto-reconnect after connection drop behaviour

Author: ntninja

tests/query_client_test.ts

@@ -13,6 +13,10 @@ import {
   assertRejects,
   assertThrows,
 } from "jsr:@std/assert@1.0.10";
+import {
+  MockError,
+  stub,
+} from "jsr:@std/testing/mock";
 import { getMainConfiguration } from "./config.ts";
 import type { PoolClient, QueryClient } from "../client.ts";
 import type { ClientOptions } from "../connection/connection_params.ts";
@@ -694,9 +698,13 @@ Deno.test(
 );
 
 // This test depends on the assumption that all clients will default to
-// one reconneciton by default
+// one reconnection by default
+//
+// It also specifically tests the case where the database system sends a
+// graceful disconnect, not the case where the database connection is just
+// closed due to the database process crashing or being killed.
 Deno.test(
-  "Default reconnection",
+  "Default reconnection after graceful database connection close",
   withClient(async (client) => {
     await assertRejects(
       () =>
@@ -714,6 +722,66 @@ Deno.test(
   }),
 );
 
+// This test depends on the assumption that all clients will default to
+// one reconnection by default
+Deno.test(
+  "Default reconnection after connection drop",
+  withClient(async (client) => {
+    // Test that connection is OK
+    const { rows: result } = await client.queryObject<{ res: number }>({
+      text: `SELECT 1`,
+      fields: ["res"],
+    });
+    assertEquals(result[0].res, 1);
+
+    // Use mock simulate irrecoverably broken pipe
+    const connWrite = stub(
+      WritableStreamDefaultWriter.prototype,
+      "write",
+      (_buffer) => {
+        throw new Deno.errors.BrokenPipe();
+      }
+    );
+    try {
+      await assertRejects(
+        () =>
+          client.queryArray`SELECT 1`,
+        Deno.errors.BrokenPipe,
+      );
+    } finally {
+      connWrite.restore();
+    }
+
+    // Use mock simulate broken pipe, that can be healed by reconnecting
+    const connWrite2 = stub(
+      WritableStreamDefaultWriter.prototype,
+      "write",
+      (_buffer) => {
+        connWrite2.restore();
+        throw new Deno.errors.BrokenPipe();
+      }
+    );
+    try {
+      const { rows: result2 } = await client.queryObject<{ res: number }>({
+        text: `SELECT 1`,
+        fields: ["res"],
+      });
+      assertEquals(result2[0].res, 1);
+    } finally {
+      // Restoring must fail here because it already happened
+      //
+      // If this fails, it means the Mock was never invoked and the test must be updated!
+      assertThrows(
+        () =>
+          connWrite2.restore(),
+        MockError,
+      );
+    }
+
+    assertEquals(client.connected, true);
+  }),
+);
+
 Deno.test(
   "Handling of debug notices",
   withClient(async (client) => {
@@ -1498,6 +1566,41 @@ Deno.test(
   }),
 );
 
+Deno.test(
+  "Transaction fails after connection drop",
+  withClient(async (client) => {
+    const name = "transactionFailsOnConnectionDrop";
+    const transaction = client.createTransaction(name);
+
+    await transaction.begin();
+
+    // Use mock simulate broken pipe
+    const connWrite = stub(
+      WritableStreamDefaultWriter.prototype,
+      "write",
+      (_buffer) => {
+        throw new Deno.errors.BrokenPipe();
+      }
+    );
+    try {
+      await assertRejects(
+        () =>
+          transaction.queryArray`SELECT 1`,
+        TransactionError,
+      );
+    } finally {
+      connWrite.restore();
+    }
+    
+    // Test that connection is OK afterward
+    const { rows: result } = await client.queryObject<{ res: number }>({
+      text: `SELECT 1`,
+      fields: ["res"],
+    });
+    assertEquals(result[0].res, 1);
+  }),
+);
+
 Deno.test(
   "Transaction savepoints",
   withClient(async (client) => {