From f01c5078353fbf4f264feef06aa5605f5bb4eeea Mon Sep 17 00:00:00 2001
From: Vladimir <sheremet.va@icloud.com>
Date: Thu, 17 Feb 2022 10:22:11 +0300
Subject: [PATCH] feat: add VITEST_WORKER_ID env

---
 docs/config/index.md                  | 4 +++-
 docs/guide/migration.md               | 6 +++++-
 packages/vitest/src/node/pool.ts      | 3 +++
 packages/vitest/src/runtime/worker.ts | 4 +++-
 packages/vitest/src/types/worker.ts   | 1 +
 test/core/test/env.test.ts            | 4 ++++
 6 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/docs/config/index.md b/docs/config/index.md
index 0b21ec128b13..467717cb17e4 100644
--- a/docs/config/index.md
+++ b/docs/config/index.md
@@ -252,7 +252,9 @@ Silent mode
 
 - **Type:** `string | string[]`
 
-Path to setup files
+Path to setup files. They will be run before each test file.
+
+You can use `process.env.VITEST_WORKER_ID` (integer-like string) inside to distinguish between threads (will always be `1`, if run with `threads: false`).
 
 ### globalSetup
 
diff --git a/docs/guide/migration.md b/docs/guide/migration.md
index 60751bd10c46..84f184c4465f 100644
--- a/docs/guide/migration.md
+++ b/docs/guide/migration.md
@@ -18,4 +18,8 @@ Unlike Jest, mocked modules in `<root>/__mocks__` are not loaded unless `vi.mock
 
 **Jasmine API**
 
-Jest exports various [`jasmine`](https://jasmine.github.io/) globals (such as `jasmine.any()`). Any such instances will need to be migrated to [their Vitest counterparts](/api/).
\ No newline at end of file
+Jest exports various [`jasmine`](https://jasmine.github.io/) globals (such as `jasmine.any()`). Any such instances will need to be migrated to [their Vitest counterparts](/api/).
+
+**Envs**
+
+Just like Jest, Vitest sets `NODE_ENV` to `test`, if it wasn't set before. Vitest also has a counterpart for `JEST_WORKER_ID` called `VITEST_WORKER_ID`, so if you rely on it, don't forget to rename it.
\ No newline at end of file
diff --git a/packages/vitest/src/node/pool.ts b/packages/vitest/src/node/pool.ts
index 7dcfd421137a..d7c5debb84d8 100644
--- a/packages/vitest/src/node/pool.ts
+++ b/packages/vitest/src/node/pool.ts
@@ -39,6 +39,7 @@ export function createFakePool(ctx: Vitest): WorkerPool {
         config: ctx.getConfig(),
         files,
         invalidates,
+        id: 1,
       }
 
       await worker[name](data, { transferList: [workerPort] })
@@ -78,6 +79,7 @@ export function createWorkerPool(ctx: Vitest): WorkerPool {
 
   const runWithFiles = (name: string): RunWithFiles => {
     return async(files, invalidates) => {
+      let id = 0
       await Promise.all(files.map(async(file) => {
         const { workerPort, port } = createChannel(ctx)
 
@@ -86,6 +88,7 @@ export function createWorkerPool(ctx: Vitest): WorkerPool {
           config: ctx.getConfig(),
           files: [file],
           invalidates,
+          id: ++id,
         }
 
         await pool.run(data, { transferList: [workerPort], name })
diff --git a/packages/vitest/src/runtime/worker.ts b/packages/vitest/src/runtime/worker.ts
index d788c16b0f9a..e2d67b6980b9 100644
--- a/packages/vitest/src/runtime/worker.ts
+++ b/packages/vitest/src/runtime/worker.ts
@@ -58,7 +58,9 @@ function init(ctx: WorkerContext) {
 
   process.stdout.write('\0')
 
-  const { config, port } = ctx
+  const { config, port, id } = ctx
+
+  process.env.VITEST_WORKER_ID = String(id)
 
   // @ts-expect-error I know what I am doing :P
   globalThis.__vitest_worker__ = {
diff --git a/packages/vitest/src/types/worker.ts b/packages/vitest/src/types/worker.ts
index 8d326adc257e..0f318363f190 100644
--- a/packages/vitest/src/types/worker.ts
+++ b/packages/vitest/src/types/worker.ts
@@ -7,6 +7,7 @@ import type { SnapshotResult } from './snapshot'
 import type { UserConsoleLog } from './general'
 
 export interface WorkerContext {
+  id: number
   port: MessagePort
   config: ResolvedConfig
   files: string[]
diff --git a/test/core/test/env.test.ts b/test/core/test/env.test.ts
index fee811c1c156..ec28127c55bc 100644
--- a/test/core/test/env.test.ts
+++ b/test/core/test/env.test.ts
@@ -17,3 +17,7 @@ test('can see env in "define"', () => {
   expect(import.meta.env.TEST_NAME).toBe('hello world')
   expect(process.env.TEST_NAME).toBe('hello world')
 })
+
+test('has worker env', () => {
+  expect(process.env.VITEST_WORKER_ID).toBeDefined()
+})