fix(pool): init VITEST_POOL_ID + VITEST_WORKER_ID before environment setup (#9085)

Author: AriPerkkio

packages/vitest/src/node/pools/pool.ts

@@ -73,6 +73,9 @@ export class Pool {
       let isMemoryLimitReached = false
       const runner = this.getPoolRunner(task, method)
 
+      const poolId = runner.poolId ?? this.getWorkerId()
+      runner.poolId = poolId
+
       const activeTask = { task, resolver, method, cancelTask }
       this.activeTasks.push(activeTask)
 
@@ -110,12 +113,9 @@ export class Pool {
           WORKER_START_TIMEOUT,
         )
 
-        await runner.start().finally(() => clearTimeout(id))
+        await runner.start({ workerId: task.context.workerId }).finally(() => clearTimeout(id))
       }
 
-      const poolId = runner.poolId ?? this.getWorkerId()
-      runner.poolId = poolId
-
       const span = runner.startTracesSpan(`vitest.worker.${method}`)
       // Start running the test in the worker
       runner.request(method, task.context)

packages/vitest/src/node/pools/poolRunner.ts

@@ -127,7 +127,6 @@ export class PoolRunner {
       __vitest_worker_request__: true,
       type: method,
       context,
-      poolId: this.poolId!,
       otelCarrier: this.getOTELCarrier(),
     } satisfies WorkerRequest)
   }
@@ -139,7 +138,7 @@ export class PoolRunner {
       : undefined
   }
 
-  async start(): Promise<void> {
+  async start(options: { workerId: number }): Promise<void> {
     // Wait for any ongoing operation to complete
     if (this._operationLock) {
       await this._operationLock
@@ -182,6 +181,8 @@ export class PoolRunner {
 
       this.postMessage({
         type: 'start',
+        poolId: this.poolId!,
+        workerId: options.workerId,
         __vitest_worker_request__: true,
         options: {
           reportMemory: this.worker.reportMemory ?? false,

packages/vitest/src/node/pools/types.ts

@@ -70,6 +70,8 @@ export type WorkerRequest
   = { __vitest_worker_request__: true } & (
     | {
       type: 'start'
+      poolId: number
+      workerId: WorkerExecuteContext['workerId'] // Initial worker ID, may change when non-isolated worker runs multiple test files
       options: { reportMemory: boolean }
       context: {
         environment: WorkerTestEnvironment
@@ -89,13 +91,11 @@ export type WorkerRequest
     | {
       type: 'run'
       context: WorkerExecuteContext
-      poolId: number
       otelCarrier?: OTELCarrier
     }
     | {
       type: 'collect'
       context: WorkerExecuteContext
-      poolId: number
       otelCarrier?: OTELCarrier
     }
     | { type: 'cancel' }

packages/vitest/src/runtime/workers/init.ts

@@ -41,6 +41,8 @@ export function init(worker: Options): void {
 
     switch (message.type) {
       case 'start': {
+        process.env.VITEST_POOL_ID = String(message.poolId)
+        process.env.VITEST_WORKER_ID = String(message.workerId)
         reportMemory = message.options.reportMemory
 
         const tracesStart = performance.now()
@@ -96,7 +98,6 @@ export function init(worker: Options): void {
         }
 
         try {
-          process.env.VITEST_POOL_ID = String(message.poolId)
           process.env.VITEST_WORKER_ID = String(message.context.workerId)
         }
         catch (error) {
@@ -155,7 +156,6 @@ export function init(worker: Options): void {
         }
 
         try {
-          process.env.VITEST_POOL_ID = String(message.poolId)
           process.env.VITEST_WORKER_ID = String(message.context.workerId)
         }
         catch (error) {

test/core/test/environments/custom-env.test.ts

@@ -6,4 +6,12 @@ test('custom env is defined', () => {
   expect(expect.getState().environment).toBe('custom')
   expect((globalThis as any).testEnvironment).toBe('custom')
   expect((globalThis as any).option).toBe('config-option')
+
+  expect((global as any).POOL_ID_DURING_ENV_SETUP).toBeDefined()
+  expect(process.env.VITEST_POOL_ID).toBeDefined()
+  expect((global as any).POOL_ID_DURING_ENV_SETUP).toBe(process.env.VITEST_POOL_ID)
+
+  expect((global as any).WORKER_ID_DURING_ENV_SETUP).toBeDefined()
+  expect(process.env.VITEST_WORKER_ID).toBeDefined()
+  expect((global as any).WORKER_ID_DURING_ENV_SETUP).toBe(process.env.VITEST_WORKER_ID)
 })

test/core/vitest-environment-custom/index.ts

@@ -12,6 +12,8 @@ export default <Environment>{
     const context = vm.createContext({
       testEnvironment: 'custom',
       option: custom.option,
+      POOL_ID_DURING_ENV_SETUP: process.env.VITEST_POOL_ID,
+      WORKER_ID_DURING_ENV_SETUP: process.env.VITEST_WORKER_ID,
       setTimeout,
       clearTimeout,
       AbortController,
@@ -31,6 +33,9 @@ export default <Environment>{
   setup(global, { custom }) {
     global.testEnvironment = 'custom'
     global.option = custom.option
+    global.POOL_ID_DURING_ENV_SETUP = process.env.VITEST_POOL_ID
+    global.WORKER_ID_DURING_ENV_SETUP = process.env.VITEST_WORKER_ID
+
     return {
       teardown() {
         delete global.testEnvironment