fix: colors on forks pool

vitest-dev · Dec 17, 2024 · 2a09cd5 · 2a09cd5
1 parent 4e60333
commit 2a09cd5
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 26 deletions.
diff --git a/packages/vitest/src/node/pool.ts b/packages/vitest/src/node/pool.ts
@@ -3,6 +3,7 @@ import type { Vitest } from './core'
 import type { TestProject } from './project'
 import type { TestSpecification } from './spec'
 import type { BuiltinPool, Pool } from './types/pool-options'
+import { isatty } from 'node:tty'
 import mm from 'micromatch'
 import { isWindows } from '../utils/env'
 import { createForksPool } from './pools/forks'
@@ -124,6 +125,7 @@ export function createPool(ctx: Vitest): ProcessPool {
         VITEST: 'true',
         NODE_ENV: process.env.NODE_ENV || 'test',
         VITEST_MODE: ctx.config.watch ? 'WATCH' : 'RUN',
+        FORCE_TTY: isatty(1) ? 'true' : '',
         ...process.env,
         ...ctx.config.env,
       },

diff --git a/test/config/test/console-color.test.ts b/test/config/test/console-color.test.ts
@@ -1,32 +1,63 @@
-import { x } from 'tinyexec'
 import { expect, test } from 'vitest'
-
-// use "tinyexec" directly since "runVitestCli" strips color
+import { runVitest } from '../../test-utils'
 
 test('with color', async () => {
-  const proc = await x('vitest', ['run', '--root=./fixtures/console-color'], {
-    nodeOptions: {
-      env: {
-        CI: '1',
-        FORCE_COLOR: '1',
-        NO_COLOR: undefined,
-        GITHUB_ACTIONS: undefined,
-      },
+  const { stdout } = await runVitest({
+    root: 'fixtures/console-color',
+    env: {
+      CI: '1',
+      FORCE_COLOR: '1',
+      NO_COLOR: undefined,
+      GITHUB_ACTIONS: undefined,
     },
-  })
-  expect(proc.stdout).toContain('\x1B[33mtrue\x1B[39m\n')
+  }, undefined, undefined, undefined, { preserveAnsi: true })
+
+  expect(stdout).toContain('\x1B[33mtrue\x1B[39m\n')
 })
 
 test('without color', async () => {
-  const proc = await x('vitest', ['run', '--root=./fixtures/console-color'], {
-    nodeOptions: {
-      env: {
-        CI: '1',
-        FORCE_COLOR: undefined,
-        NO_COLOR: '1',
-        GITHUB_ACTIONS: undefined,
-      },
+  const { stdout } = await runVitest({
+    root: 'fixtures/console-color',
+    env: {
+      CI: '1',
+      FORCE_COLOR: undefined,
+      NO_COLOR: '1',
+      GITHUB_ACTIONS: undefined,
+    },
+  }, undefined, undefined, undefined, { preserveAnsi: true })
+
+  expect(stdout).toContain('true\n')
+  expect(stdout).not.toContain('\x1B[33mtrue\x1B[39m\n')
+})
+
+test('without color, forks pool in non-TTY parent', async () => {
+  const { stdout } = await runVitest({
+    root: 'fixtures/console-color',
+    env: {
+      CI: undefined,
+      FORCE_COLOR: undefined,
+      NO_COLOR: undefined,
+      GITHUB_ACTIONS: undefined,
+
+      // Overrides current process's value, since we are running Vitest in Vitest here
+      FORCE_TTY: undefined,
+    },
+  }, undefined, undefined, undefined, { preserveAnsi: true })
+
+  expect(stdout).toContain('true\n')
+  expect(stdout).not.toContain('\x1B[33mtrue\x1B[39m\n')
+})
+
+test('with color, forks pool in TTY parent', async () => {
+  const { stdout } = await runVitest({
+    root: 'fixtures/console-color',
+    env: {
+      CI: undefined,
+      FORCE_COLOR: undefined,
+      NO_COLOR: undefined,
+      GITHUB_ACTIONS: undefined,
     },
-  })
-  expect(proc.stdout).toContain('true\n')
+  }, undefined, undefined, undefined, { preserveAnsi: true })
+
+  expect(stdout).toContain('\x1B[33mtrue\x1B[39m\n')
 })
diff --git a/test/test-utils/cli.ts b/test/test-utils/cli.ts
@@ -12,9 +12,11 @@ export class Cli {
   private stdoutListeners: Listener[] = []
   private stderrListeners: Listener[] = []
   private stdin: ReadableOrWritable
+  private preserveAnsi?: boolean
 
-  constructor(options: { stdin: ReadableOrWritable; stdout: ReadableOrWritable; stderr: ReadableOrWritable }) {
+  constructor(options: { stdin: ReadableOrWritable; stdout: ReadableOrWritable; stderr: ReadableOrWritable; preserveAnsi?: boolean }) {
     this.stdin = options.stdin
+    this.preserveAnsi = options.preserveAnsi
 
     for (const source of (['stdout', 'stderr'] as const)) {
       const stream = options[source]
@@ -37,7 +39,7 @@ export class Cli {
   }
 
   private capture(source: Source, data: any) {
-    const msg = stripVTControlCharacters(data.toString())
+    const msg = this.preserveAnsi ? data.toString() : stripVTControlCharacters(data.toString())
     this[source] += msg
     this[`${source}Listeners`].forEach(fn => fn())
   }

diff --git a/test/test-utils/index.ts b/test/test-utils/index.ts
@@ -20,6 +20,7 @@ Object.assign(tinyrainbow.default, tinyrainbow.getDefaultColors())
 interface VitestRunnerCLIOptions {
   std?: 'inherit'
   fails?: boolean
+  preserveAnsi?: boolean
 }
 
 export async function runVitest(
@@ -58,7 +59,7 @@ export async function runVitest(
   const stdin = new Readable({ read: () => '' }) as NodeJS.ReadStream
   stdin.isTTY = true
   stdin.setRawMode = () => stdin
-  const cli = new Cli({ stdin, stdout, stderr })
+  const cli = new Cli({ stdin, stdout, stderr, preserveAnsi: runnerOptions.preserveAnsi })
 
   let ctx: Vitest | undefined
   let thrown = false