From 70bf7d663cc3b8274b6be54d2162fe859d3cfc42 Mon Sep 17 00:00:00 2001 From: jxom <7336481+jxom@users.noreply.github.com> Date: Sat, 18 Jul 2026 19:33:20 +1000 Subject: [PATCH] feat: add Vitest pool utilities --- .changeset/fresh-pools-smile.md | 5 + README.md | 52 ++++++++- package.json | 5 + src/Instance.test.ts | 4 + src/Instance.ts | 6 + src/Server.test.ts | 3 + src/Server.ts | 3 +- src/vitest/Pool.test.ts | 198 ++++++++++++++++++++++++++++++++ src/vitest/Pool.ts | 92 +++++++++++++++ src/vitest/index.ts | 1 + tsconfig.json | 3 +- 11 files changed, 369 insertions(+), 3 deletions(-) create mode 100644 .changeset/fresh-pools-smile.md create mode 100644 src/vitest/Pool.test.ts create mode 100644 src/vitest/Pool.ts create mode 100644 src/vitest/index.ts diff --git a/.changeset/fresh-pools-smile.md b/.changeset/fresh-pools-smile.md new file mode 100644 index 0000000..f6c7a0c --- /dev/null +++ b/.changeset/fresh-pools-smile.md @@ -0,0 +1,5 @@ +--- +'prool': patch +--- + +Added Vitest pool helpers and a derived `Instance.url` property. diff --git a/README.md b/README.md index 7ea09f9..0597b44 100644 --- a/README.md +++ b/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)) diff --git a/package.json b/package.json index 88632de..5bca48b 100644 --- a/package.json +++ b/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": { diff --git a/src/Instance.test.ts b/src/Instance.test.ts index 0034386..5b3de0b 100644 --- a/src/Instance.test.ts +++ b/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({ diff --git a/src/Instance.ts b/src/Instance.ts index 7f90606..7c1aa71 100644 --- a/src/Instance.ts +++ b/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, get status() { if (restarting) return 'restarting' diff --git a/src/Server.test.ts b/src/Server.test.ts index 26c4da7..b2efc5e 100644 --- a/src/Server.test.ts +++ b/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, diff --git a/src/Server.ts b/src/Server.ts index a44740a..90ed47f 100644 --- a/src/Server.ts +++ b/src/Server.ts @@ -319,12 +319,13 @@ function websocketProxyTarget(endpoint: Endpoint) { } function instanceDescriptor( - instance: Pick, + instance: Pick, ) { return { endpoints: instance.endpoints, host: instance.host, port: instance.port, + url: instance.url, } } diff --git a/src/vitest/Pool.test.ts b/src/vitest/Pool.test.ts new file mode 100644 index 0000000..a9bc929 --- /dev/null +++ b/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() + return { + context, + project: { + config: { maxWorkers }, + provide(key: string, value: unknown) { + context.set(key, value) + }, + }, + } +} diff --git a/src/vitest/Pool.ts b/src/vitest/Pool.ts new file mode 100644 index 0000000..a3809d6 --- /dev/null +++ b/src/vitest/Pool.ts @@ -0,0 +1,92 @@ +import type { Instance } from '../Instance.js' +import * as Pool from '../Pool.js' + +type StartedInstance = Awaited< + ReturnType['start']> +> + +/** Returns the value provided for the current Vitest pool. */ +export function get(values: readonly value[]): value { + const id = poolId() + const index = id - 1 + if (index >= values.length) + throw new Error(`Missing value for Vitest pool ${id}.`) + return values[index]! +} + +/** Returns the current Vitest pool ID. */ +export function poolId(): number { + const value = process.env['VITEST_POOL_ID'] + if (value === undefined) throw new Error('VITEST_POOL_ID is not set.') + + const id = Number(value) + if (!Number.isSafeInteger(id) || id < 1) + throw new Error(`Invalid VITEST_POOL_ID "${value}".`) + return id +} + +/** Creates Vitest global setup with one instance per worker. */ +export function setup< + instance extends Instance = Instance, + project extends setup.Project = setup.Project, +>(parameters: setup.Parameters): setup.ReturnType { + return async (project) => { + const { maxWorkers } = project.config + if (!Number.isSafeInteger(maxWorkers) || maxWorkers < 1) + throw new Error('Vitest maxWorkers must be a positive integer.') + + const pool = Pool.define({ + instance: parameters.instance, + limit: maxWorkers, + }) + const starts = Array.from({ length: maxWorkers }, (_, index) => + pool.start(index + 1), + ) + + try { + const instances = await Promise.all(starts) + await parameters.setup(instances, project) + return () => pool.destroyAll() + } catch (error) { + await Promise.allSettled(starts) + try { + await pool.destroyAll() + } catch (destroyError) { + throw new AggregateError( + [error, destroyError], + 'Failed to set up or destroy Vitest pool.', + ) + } + throw error + } + } +} + +export declare namespace setup { + /** Options for setting up a Vitest worker pool. */ + export type Parameters< + instance extends Instance = Instance, + project extends Project = Project, + > = { + /** Instance created for each Vitest pool ID. */ + instance: instance | ((poolId: number) => instance) + /** Configures serializable context provided to every worker. */ + setup( + instances: readonly StartedInstance[], + project: project, + ): Promise | void + } + + /** Minimal Vitest project interface used by global setup. */ + export type Project = { + config: { + maxWorkers: number + } + provide(key: string, value: unknown): void + } + + /** Vitest global setup function. */ + export type ReturnType = ( + project: project, + ) => Promise<() => Promise> +} diff --git a/src/vitest/index.ts b/src/vitest/index.ts new file mode 100644 index 0000000..17f8eb1 --- /dev/null +++ b/src/vitest/index.ts @@ -0,0 +1 @@ +export * as Pool from './Pool.js' diff --git a/tsconfig.json b/tsconfig.json index 71aae58..9495dd5 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -8,7 +8,8 @@ "target": "esnext", "paths": { "prool": ["./src/index.ts"], - "prool/testcontainers": ["./src/testcontainers/index.ts"] + "prool/testcontainers": ["./src/testcontainers/index.ts"], + "prool/vitest": ["./src/vitest/index.ts"] }, "types": ["node", "vitest/globals"],