From 8d2a5ea7d93613d6ca9c14017f29ab092fa1327e Mon Sep 17 00:00:00 2001
From: Chapman Pendery <cpendery@microsoft.com>
Date: Wed, 6 Mar 2024 12:05:01 -0800
Subject: [PATCH] feat: make worker count configurable

Signed-off-by: Chapman Pendery <cpendery@microsoft.com>
---
 src/config/config.ts | 23 +++++++++++++++++++++++
 src/reporter/base.ts |  7 +++++--
 src/reporter/list.ts |  8 ++++++--
 src/runner/runner.ts | 37 +++++++++++++++++++++----------------
 4 files changed, 55 insertions(+), 20 deletions(-)

diff --git a/src/config/config.ts b/src/config/config.ts
index 9df09c0..cf1813d 100644
--- a/src/config/config.ts
+++ b/src/config/config.ts
@@ -4,6 +4,7 @@
 import fs from "node:fs";
 import path from "node:path";
 import process from "node:process";
+import os from "node:os";
 
 import { defaultShell } from "../terminal/shell.js";
 import { TestOptions } from "../test/option.js";
@@ -26,6 +27,10 @@ export const loadConfig = async (): Promise<Required<TestConfig>> => {
     projects: userConfig.projects ?? [],
     timeout: userConfig.timeout ?? 30_000,
     reporter: userConfig.reporter ?? "list",
+    workers: Math.max(
+      userConfig.workers ?? Math.max(Math.floor(os.cpus().length / 2), 1),
+      1
+    ),
     use: {
       shell: userConfig.use?.shell ?? defaultShell,
       rows: userConfig.use?.rows ?? 30,
@@ -197,6 +202,24 @@ export declare type TestConfig = {
    */
   use?: TestOptions;
 
+  /**
+   * The number of workers to use. Defaults to 50% of the logical cpu cores. If
+   * there are less tests than requested workers, there will be 1 worker used per test.
+   *
+   * **Usage**
+   *
+   * ```js
+   * // tui-test.config.ts
+   * import { defineConfig } from '@microsoft/tui-test';
+   *
+   * export default defineConfig({
+   *   workers: 4,
+   * });
+   * ```
+   *
+   */
+  workers?: number;
+
   /**
    * TUI Test supports running multiple test projects at the same time.
    *
diff --git a/src/reporter/base.ts b/src/reporter/base.ts
index cf45195..946a977 100644
--- a/src/reporter/base.ts
+++ b/src/reporter/base.ts
@@ -7,7 +7,6 @@ import { TestCase, TestResult, TestStatus } from "../test/testcase.js";
 import { Shell } from "../terminal/shell.js";
 import { loadShellVersions } from "./utils.js";
 import { Suite } from "../test/suite.js";
-import { maxWorkers } from "../runner/runner.js";
 
 type TestSummary = {
   didNotRun: number;
@@ -44,7 +43,11 @@ export class BaseReporter {
     return `${text}${count > 1 ? "s" : ""}`;
   }
 
-  async start(testCount: number, shells: Shell[]): Promise<void> {
+  async start(
+    testCount: number,
+    shells: Shell[],
+    maxWorkers: number
+  ): Promise<void> {
     const shellVersions = await loadShellVersions(shells);
     const workers = Math.min(testCount, maxWorkers);
     process.stdout.write(
diff --git a/src/reporter/list.ts b/src/reporter/list.ts
index 903367e..75abda9 100644
--- a/src/reporter/list.ts
+++ b/src/reporter/list.ts
@@ -18,8 +18,12 @@ export class ListReporter extends BaseReporter {
     this._testRows = {};
   }
 
-  override async start(testCount: number, shells: Shell[]): Promise<void> {
-    await super.start(testCount, shells);
+  override async start(
+    testCount: number,
+    shells: Shell[],
+    maxWorkers: number
+  ): Promise<void> {
+    await super.start(testCount, shells, maxWorkers);
   }
 
   override startTest(test: TestCase, result: TestResult): void {
diff --git a/src/runner/runner.ts b/src/runner/runner.ts
index 304b094..ce6ddc9 100644
--- a/src/runner/runner.ts
+++ b/src/runner/runner.ts
@@ -4,7 +4,7 @@
 import path from "node:path";
 import url from "node:url";
 import os from "node:os";
-import workerpool from "workerpool";
+import workerpool, { Pool } from "workerpool";
 import chalk from "chalk";
 
 import { Suite, getRootSuite } from "../test/suite.js";
@@ -32,25 +32,28 @@ type ExecutionOptions = {
 };
 
 const __dirname = path.dirname(url.fileURLToPath(import.meta.url));
-export const maxWorkers = Math.max(Math.floor(os.cpus().length / 2), 1);
-const pool = workerpool.pool(path.join(__dirname, "worker.js"), {
-  workerType: "process",
-  maxWorkers,
-  forkOpts: {
-    stdio: "inherit",
-    env: {
-      ...process.env,
-      ...(supportsColor ? { FORCE_COLOR: "1" } : {}),
+
+const createWorkerPool = (maxWorkers: number): Pool => {
+  return workerpool.pool(path.join(__dirname, "worker.js"), {
+    workerType: "process",
+    maxWorkers,
+    forkOpts: {
+      stdio: "inherit",
+      env: {
+        ...process.env,
+        ...(supportsColor ? { FORCE_COLOR: "1" } : {}),
+      },
     },
-  },
-  emitStdStreams: true,
-});
+    emitStdStreams: true,
+  });
+};
 
 const runSuites = async (
   allSuites: Suite[],
   filteredTestIds: Set<string>,
   reporter: BaseReporter,
-  { updateSnapshot }: ExecutionOptions
+  { updateSnapshot }: ExecutionOptions,
+  pool: Pool
 ) => {
   const tasks: Promise<void>[] = [];
   const suites = [...allSuites];
@@ -161,6 +164,7 @@ export const run = async (options: ExecutionOptions) => {
   const config = await loadConfig();
   const rootSuite = await getRootSuite(config);
   const reporter = new ListReporter();
+  const pool = createWorkerPool(config.workers);
 
   const suites = [rootSuite];
   while (suites.length != 0) {
@@ -217,7 +221,7 @@ export const run = async (options: ExecutionOptions) => {
   if (shells.includes(Shell.Zsh)) {
     await setupZshDotfiles();
   }
-  await reporter.start(allTests.length, shells);
+  await reporter.start(allTests.length, shells, config.workers);
 
   if (config.globalTimeout > 0) {
     setTimeout(() => {
@@ -232,7 +236,8 @@ export const run = async (options: ExecutionOptions) => {
     rootSuite.suites,
     new Set(allTests.map((test) => test.id)),
     reporter,
-    options
+    options,
+    pool
   );
   try {
     await pool.terminate(true);