Cleanup

Author: codetheweb

src/index.ts

@@ -4,6 +4,8 @@ import path from "node:path"
 import type {
   ConnectionDetailsFromWorker,
   InitialWorkerData,
+  SharedWorkerFunctions,
+  TestWorkerFunctions,
 } from "./internal-types"
 import type {
   ConnectionDetails,
@@ -16,7 +18,6 @@ import type { Jsonifiable } from "type-fest"
 import type { ExecutionContext } from "ava"
 import { once } from "node:events"
 import { createBirpc } from "birpc"
-import { SharedWorkerFunctions, TestWorkerFunctions } from "./lib/rpc"
 import { ExecResult } from "testcontainers"
 import isPlainObject from "lodash/isPlainObject"
 
@@ -187,23 +188,19 @@ export const getTestPostgresDatabaseFactory = <
     }
   )
 
-  // todo: properly tear down?
-  const messageHandlerAbortController = new AbortController()
-  const messageHandlerPromise = Promise.race([
-    once(messageHandlerAbortController.signal, "abort"),
-    (async () => {
-      const worker = await workerPromise
-      await worker.available
+  // Automatically cleaned up by AVA since each test file runs in a separate worker
+  const _messageHandlerPromise = (async () => {
+    const worker = await workerPromise
+    await worker.available
 
-      for await (const msg of worker.subscribe()) {
-        rpcCallback!(msg.data)
-
-        if (messageHandlerAbortController.signal.aborted) {
-          break
-        }
+    for await (const msg of worker.subscribe()) {
+      if (msg.data.type === "teardown") {
+        console.log("breaking")
+        break
       }
-    })(),
-  ])
+      rpcCallback!(msg.data)
+    }
+  })()
 
   const getTestPostgresDatabase: GetTestPostgresDatabase<Params> = async (
     t: ExecutionContext,

src/internal-types.ts

@@ -1,5 +1,6 @@
-import type { ExecResult, StartedNetwork } from "testcontainers"
+import type { ExecResult } from "testcontainers"
 import type { Jsonifiable } from "type-fest"
+
 import type {
   ConnectionDetails,
   GetTestPostgresDatabaseFactoryOptions,
@@ -13,52 +14,20 @@ export interface InitialWorkerData {
 
 export type ConnectionDetailsFromWorker = Omit<ConnectionDetails, "pool">
 
-export interface RequestDatabaseFromWorkerMessage {
-  type: "GET_TEST_DATABASE"
-  key?: string
-  params?: Jsonifiable
-}
-
-export interface RequestBeforeTemplateIsBakedHookToBeRunMessage {
-  type: "RUN_HOOK_BEFORE_TEMPLATE_IS_BAKED"
-  connectionDetails: ConnectionDetailsFromWorker
-}
-
-export interface FinishedRunningBeforeTemplateIsBakedHookMessage {
-  type: "FINISHED_RUNNING_HOOK_BEFORE_TEMPLATE_IS_BAKED"
-  result:
-    | {
-        status: "success"
-        result: any
-      }
-    | {
-        status: "error"
-        error: Error | string
-      }
+export interface TestWorkerFunctions {
+  runBeforeTemplateIsBakedHook: (
+    connectionDetails: ConnectionDetailsFromWorker,
+    params?: Jsonifiable
+  ) => Promise<unknown>
 }
 
-export interface ExecCommandInContainerMessage {
-  type: "EXEC_COMMAND_IN_CONTAINER"
-  command: string[]
+export interface SharedWorkerFunctions {
+  getTestDatabase: (options: {
+    key?: string
+    params?: Jsonifiable
+  }) => Promise<{
+    connectionDetails: ConnectionDetailsFromWorker
+    beforeTemplateIsBakedResult: unknown
+  }>
+  execCommandInContainer: (command: string[]) => Promise<ExecResult>
 }
-
-export interface ExecCommandInContainerResultMessage {
-  type: "EXEC_COMMAND_IN_CONTAINER_RESULT"
-  result: ExecResult
-}
-
-export interface GotDatabaseMessage {
-  type: "GOT_DATABASE"
-  connectionDetails: ConnectionDetailsFromWorker
-  beforeTemplateIsBakedResult: FinishedRunningBeforeTemplateIsBakedHookMessage["result"]
-}
-
-export type MessageToWorker =
-  | RequestDatabaseFromWorkerMessage
-  | FinishedRunningBeforeTemplateIsBakedHookMessage
-  | ExecCommandInContainerMessage
-export type MessageFromWorker =
-  | RequestBeforeTemplateIsBakedHookToBeRunMessage
-  | GotDatabaseMessage
-  | ExecCommandInContainerResultMessage
-export type WorkerMessage = MessageToWorker | MessageFromWorker

src/lib/rpc.ts

@@ -2,12 +2,15 @@ import pg from "pg"
 import { GenericContainer, Network } from "testcontainers"
 import { Mutex } from "async-mutex"
 import hash from "object-hash"
-import type { InitialWorkerData } from "./internal-types"
+import type {
+  InitialWorkerData,
+  SharedWorkerFunctions,
+  TestWorkerFunctions,
+} from "./internal-types"
 import getRandomDatabaseName from "./lib/get-random-database-name"
 import type { SharedWorker } from "ava/plugin"
 import { type BirpcReturn, type ChannelOptions, createBirpc } from "birpc"
 import { once } from "node:events"
-import type { SharedWorkerFunctions, TestWorkerFunctions } from "./lib/rpc"
 import type { Jsonifiable } from "type-fest"
 
 type WorkerRpc = BirpcReturn<TestWorkerFunctions, SharedWorkerFunctions>
@@ -54,6 +57,7 @@ export class Worker {
     testWorker.teardown(async () => {
       messageHandlerAbortController.abort()
       await messageHandlerPromise
+      testWorker.publish({ type: "teardown" })
     })
 
     const rpc: WorkerRpc = createBirpc<