feat: add Vitest pool utilities

.changeset/fresh-pools-smile.md

@@ -0,0 +1,5 @@
+---
+'prool': patch
+---
+
+Added Vitest pool helpers and a derived `Instance.url` property.

README.md

@@ -54,6 +54,7 @@ You can also create your own custom instances by using the [`Instance.define` fu
   - [`Instance.define`](#instancedefine)
   - [`Pool.create`](#poolcreate)
   - [`Pool.define`](#pooldefine)
+  - [Vitest](#vitest)
 
 
 ## Install
@@ -266,7 +267,7 @@ const pool = Pool.create({
 })
 const lease = await pool.acquire()
 try {
-  await fetch(`http://${lease.instance.host}:${lease.instance.port}`)
+  await fetch(lease.instance.url)
 } finally {
   await lease.release()
 }
@@ -307,6 +308,55 @@ const instance_3 = await pool.start(3)
 | `limit`    | Number of instances that can be instantiated in the pool | `number`   |
 | returns    | Pool.                                                    | `Pool`     |
 
+### Vitest
+
+`Pool.setup` starts one instance per Vitest Node worker and runs `setup` with the instances and Vitest project. `Pool.get` selects the current worker's value from a provided array. Configure the global setup on each project that uses the instances.
+
+#### Config
+
+```ts
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    globalSetup: './test/setup.global.ts',
+  },
+})
+```
+
+#### Global setup
+
+```ts
+import { Instance } from 'prool'
+import { Pool } from 'prool/vitest'
+import type { TestProject } from 'vitest/node'
+
+declare module 'vitest' {
+  export interface ProvidedContext {
+    rpcUrls: readonly string[]
+  }
+}
+
+export default Pool.setup({
+  instance: Instance.anvil(),
+  setup(instances, project: TestProject) {
+    project.provide(
+      'rpcUrls',
+      instances.map((instance) => instance.url),
+    )
+  },
+})
+```
+
+#### Worker
+
+```ts
+import { Pool } from 'prool/vitest'
+import { inject } from 'vitest'
+
+export const rpcUrl = Pool.get(inject('rpcUrls'))
+```
+
 ## Authors
 
 - [@jxom](https://github.com/jxom) (jxom.eth, [Twitter](https://twitter.com/jakemoxey))

package.json

@@ -48,6 +48,11 @@
       "src": "./src/testcontainers/index.ts",
       "types": "./dist/testcontainers/index.d.ts",
       "default": "./dist/testcontainers/index.js"
+    },
+    "./vitest": {
+      "src": "./src/vitest/index.ts",
+      "types": "./dist/vitest/index.d.ts",
+      "default": "./dist/vitest/index.js"
     }
   },
   "dependencies": {

src/Instance.test.ts

@@ -321,10 +321,12 @@ test('behavior: dynamic host/port via setEndpoint', async () => {
   const instance = foo()
   expect(instance.host).toEqual('localhost')
   expect(instance.port).toEqual(3000)
+  expect(instance.url).toEqual('http://localhost:3000')
 
   await instance.start()
   expect(instance.host).toEqual('192.168.1.100')
   expect(instance.port).toEqual(9999)
+  expect(instance.url).toEqual('http://192.168.1.100:9999')
 })
 
 test('behavior: named endpoints', async () => {
@@ -387,6 +389,7 @@ test('behavior: named endpoints', async () => {
     port: 3100,
     protocol: 'https',
   })
+  expect(instance.url).toBe('https://endpoint.localhost:3100')
   expect(instance.endpoints.metrics).toEqual({
     host: 'localhost',
     port: 9090,
@@ -398,6 +401,7 @@ test('behavior: named endpoints', async () => {
 
   expect(instance.host).toEqual('127.0.0.1')
   expect(instance.port).toEqual(4000)
+  expect(instance.url).toEqual('https://127.0.0.1:4000')
   expect(instance.endpoints.default.protocol).toBe('https')
   expect(instance.endpoints.default).toBe(endpoints.default)
   expect(instance.endpoints.metrics).toEqual({

src/Instance.ts

@@ -71,6 +71,8 @@ export type Instance<
    * Port the instance is running on.
    */
   port: number
+  /** URL of the default endpoint. */
+  url: string
   /**
    * Set of messages emitted from the `"message"` event stored in-memory,
    * with length {@link InstanceOptions`messageBuffer`}.
@@ -221,6 +223,10 @@ export function define<
         get port() {
           return endpointMap.default.port
         },
+        get url() {
+          const { host, port, protocol } = endpointMap.default
+          return `${protocol}://${host}:${port}`
+        },
         endpoints: endpointMap as InstanceEndpoints<endpoints>,
         get status() {
           if (restarting) return 'restarting'

src/Server.test.ts

@@ -62,6 +62,7 @@ test('request: lifecycle endpoint discovery', async () => {
   expect(start).toMatchObject({
     host: start.endpoints.default.host,
     port: start.endpoints.default.port,
+    url: `http://${start.endpoints.default.host}:${start.endpoints.default.port}`,
   })
 
   const restart = await fetch(`http://localhost:${port}/1/restart`).then(
@@ -110,6 +111,7 @@ test('request: leases pooled instances', async () => {
     },
     host: '127.0.0.1',
     token: expect.any(String),
+    url: `http://127.0.0.1:${first.port}`,
   })
 
   const waiting = fetch(`${url}/acquire`, { method: 'POST' })
@@ -308,6 +310,7 @@ describe.each([
     const json = (await response.json()) as any
     expect(json.host).toBeDefined()
     expect(json.port).toBeDefined()
+    expect(json.url).toBe(`http://${json.host}:${json.port}`)
     expect(json.endpoints.default).toEqual({
       host: json.host,
       port: json.port,

src/Server.ts

@@ -319,12 +319,13 @@ function websocketProxyTarget(endpoint: Endpoint) {
 }
 
 function instanceDescriptor(
-  instance: Pick<Instance, 'endpoints' | 'host' | 'port'>,
+  instance: Pick<Instance, 'endpoints' | 'host' | 'port' | 'url'>,
 ) {
   return {
     endpoints: instance.endpoints,
     host: instance.host,
     port: instance.port,
+    url: instance.url,
   }
 }
 

src/vitest/Pool.test.ts

@@ -0,0 +1,198 @@
+import { Instance } from 'prool'
+import { Pool } from 'prool/vitest'
+import { afterEach, describe, expect, expectTypeOf, test, vi } from 'vitest'
+
+afterEach(() => {
+  vi.unstubAllEnvs()
+})
+
+describe('get', () => {
+  test('returns the value for the current pool', () => {
+    vi.stubEnv('VITEST_POOL_ID', '2')
+
+    expect(Pool.get(['first', 'second', 'third'])).toBe('second')
+  })
+
+  test('throws when the current pool has no value', () => {
+    vi.stubEnv('VITEST_POOL_ID', '3')
+
+    expect(() => Pool.get(['first', 'second'])).toThrowError(
+      'Missing value for Vitest pool 3.',
+    )
+  })
+})
+
+describe('poolId', () => {
+  test('returns the current pool ID', () => {
+    vi.stubEnv('VITEST_POOL_ID', '3')
+
+    expect(Pool.poolId()).toBe(3)
+  })
+
+  test.each(['0', '-1', '1.5', 'worker'])('rejects %s', (value) => {
+    vi.stubEnv('VITEST_POOL_ID', value)
+
+    expect(() => Pool.poolId()).toThrowError(
+      `Invalid VITEST_POOL_ID "${value}".`,
+    )
+  })
+
+  test('requires VITEST_POOL_ID', () => {
+    vi.stubEnv('VITEST_POOL_ID', undefined)
+
+    expect(() => Pool.poolId()).toThrowError('VITEST_POOL_ID is not set.')
+  })
+})
+
+describe('setup', () => {
+  test('starts one instance per worker and provides setup context', async () => {
+    const started: number[] = []
+    const stopped: number[] = []
+    const { context, project } = testProject(3)
+    const setup = Pool.setup({
+      instance: (id) =>
+        Instance.define(() => ({
+          endpoints: {
+            metrics: {
+              host: 'localhost',
+              port: 9000 + id,
+              protocol: 'http' as const,
+            },
+          },
+          host: 'localhost',
+          name: `worker-${id}`,
+          port: 3000 + id,
+          async start() {
+            started.push(id)
+          },
+          async stop() {
+            stopped.push(id)
+          },
+        }))(),
+      setup(instances, project) {
+        expectTypeOf(
+          instances[0]!.endpoints.metrics.protocol,
+        ).toEqualTypeOf<'http'>()
+        project.provide(
+          'names',
+          instances.map((instance) => instance.name),
+        )
+        project.provide(
+          'urls',
+          instances.map((instance) => instance.url),
+        )
+      },
+    })
+
+    const teardown = await setup(project)
+
+    expect(started).toEqual([1, 2, 3])
+    expect(context.get('names')).toEqual(['worker-1', 'worker-2', 'worker-3'])
+    expect(context.get('urls')).toEqual([
+      expect.stringMatching(/^http:\/\/localhost:\d+$/),
+      expect.stringMatching(/^http:\/\/localhost:\d+$/),
+      expect.stringMatching(/^http:\/\/localhost:\d+$/),
+    ])
+    await teardown()
+    expect(stopped).toEqual([1, 2, 3])
+  })
+
+  test('destroys instances when setup fails', async () => {
+    const stopped: number[] = []
+    const setup = Pool.setup({
+      instance: (id) =>
+        Instance.define(() => ({
+          host: 'localhost',
+          name: `worker-${id}`,
+          port: 3000 + id,
+          async start() {},
+          async stop() {
+            stopped.push(id)
+          },
+        }))(),
+      setup() {
+        throw new Error('setup failed')
+      },
+    })
+
+    await expect(setup(testProject(2).project)).rejects.toThrowError(
+      'setup failed',
+    )
+
+    expect(stopped).toEqual([1, 2])
+  })
+
+  test('destroys started instances when a start fails', async () => {
+    const stopped: number[] = []
+    const setup = Pool.setup({
+      instance: (id) =>
+        Instance.define(() => ({
+          host: 'localhost',
+          name: `worker-${id}`,
+          port: 3000 + id,
+          async start() {
+            if (id === 2) throw new Error('start failed')
+          },
+          async stop() {
+            stopped.push(id)
+          },
+        }))(),
+      setup() {},
+    })
+
+    await expect(setup(testProject(3).project)).rejects.toThrowError(
+      'start failed',
+    )
+
+    expect(stopped).toEqual([1, 3])
+  })
+
+  test('reports setup and teardown failures', async () => {
+    const setup = Pool.setup({
+      instance: Instance.define(() => ({
+        host: 'localhost',
+        name: 'worker',
+        port: 3000,
+        async start() {},
+        async stop() {
+          throw new Error('stop failed')
+        },
+      }))(),
+      setup() {
+        throw new Error('setup failed')
+      },
+    })
+
+    const error = await setup(testProject(1).project).catch((error) => error)
+
+    expect(error).toBeInstanceOf(AggregateError)
+    expect(error.errors.map((error: Error) => error.message)).toEqual([
+      'setup failed',
+      'stop failed',
+    ])
+  })
+
+  test('requires a positive worker count', async () => {
+    const setup = Pool.setup({
+      instance: Instance.anvil(),
+      setup() {},
+    })
+
+    await expect(setup(testProject(0).project)).rejects.toThrowError(
+      'Vitest maxWorkers must be a positive integer.',
+    )
+  })
+})
+
+function testProject(maxWorkers: number) {
+  const context = new Map<string, unknown>()
+  return {
+    context,
+    project: {
+      config: { maxWorkers },
+      provide(key: string, value: unknown) {
+        context.set(key, value)
+      },
+    },
+  }
+}