feat: reuse worker process between test types (#242)

When two test types only differ in the last few environments, and none of them has `beforeAll` or `afterAll`, we can reuse the worker. In this case, the worker won't have to run new `beforeAll`/`afterAll` methods, and will update the list of envs to run `beforeEach`/`afterEach` for the next test.
microsoft · Apr 23, 2021 · 9484a63 · 9484a63
1 parent 66de782
commit 9484a63
Show file tree

Hide file tree

Showing 6 changed files with 114 additions and 14 deletions.
diff --git a/src/index.ts b/src/index.ts
@@ -26,7 +26,7 @@ import ListReporter from './reporters/list';
 export * from './types';
 export { expect } from './expect';
 export { setConfig, setReporters, globalSetup, globalTeardown } from './spec';
-export const test: TestType<{}, {}, {}, {}> = newTestTypeImpl([]);
+export const test: TestType<{}, {}, {}, {}> = newTestTypeImpl([], undefined);
 export const reporters = {
   dot: DotReporter,
   json: JSONReporter,

diff --git a/src/loader.ts b/src/loader.ts
@@ -90,14 +90,32 @@ export class Loader {
     const result = new Set<{
       fileSuites: Map<string, Suite>;
       envs: Env<any>[];
+      envHash: string;
     }>();
-    const visit = (t: TestType<any, any, any, any>) => {
+    type AnyTestType = TestType<any, any, any, any>;
+    const hashByTestType = new Map<AnyTestType, string>();
+
+    const visit = (t: AnyTestType, lastWithForkingEnv: AnyTestType) => {
       const description = configFile.testTypeDescriptions.get(t)!;
-      result.add({ fileSuites: description.fileSuites, envs: description.envs });
+
+      // Fork if we get an environment with worker-level hooks.
+      if (description.newEnv && (description.newEnv.beforeAll || description.newEnv.afterAll))
+        lastWithForkingEnv = t;
+      let envHash = hashByTestType.get(lastWithForkingEnv);
+      if (!envHash) {
+        envHash = String(hashByTestType.size);
+        hashByTestType.set(lastWithForkingEnv, envHash);
+      }
+
+      result.add({
+        fileSuites: description.fileSuites,
+        envs: description.envs,
+        envHash
+      });
       for (const child of description.children)
-        visit(child);
+        visit(child, lastWithForkingEnv);
     };
-    visit(runList.testType);
+    visit(runList.testType, runList.testType);
     return result;
   }
 

diff --git a/src/runner.ts b/src/runner.ts
@@ -57,12 +57,10 @@ export class Runner {
     const grepMatcher = createMatcher(this._loader.config().grep);
 
     const nonEmptySuites = new Set<Suite>();
-    let descriptionIndex = 0;
     for (const runList of this._loader.runLists()) {
       if (tagFilter && !runList.tags.some(tag => tagFilter.includes(tag)))
         continue;
       for (const description of this._loader.descriptionsForRunList(runList)) {
-        ++descriptionIndex;
         for (const fileSuite of description.fileSuites.values()) {
           if (filtered.size && !filtered.has(fileSuite))
             continue;
@@ -78,7 +76,7 @@ export class Runner {
                 outputDir: config.outputDir,
                 repeatEachIndex: i,
                 runListIndex: runList.index,
-                workerHash: `${descriptionIndex}#repeat-${i}`,
+                workerHash: `#list-${runList.index}#env-${description.envHash}#repeat-${i}`,
                 variationId: `#run-${runList.index}#repeat-${i}`,
               };
               spec._appendTest(testVariation);

diff --git a/src/spec.ts b/src/spec.ts
@@ -32,6 +32,7 @@ export type TestTypeDescription = {
   fileSuites: Map<string, Suite>;
   children: Set<AnyTestType>;
   envs: Env<any>[];
+  newEnv: Env<any> | undefined;
 };
 export type RunListDescription = {
   index: number;
@@ -63,12 +64,13 @@ export function setLoadingConfigFile(loading: boolean) {
 
 const countByFile = new Map<string, number>();
 
-export function newTestTypeImpl(envs: Env<any>[]): any {
+export function newTestTypeImpl(envs: Env<any>[], newEnv: Env<any> | undefined): any {
   const fileSuites = new Map<string, Suite>();
   const description: TestTypeDescription = {
     fileSuites,
     children: new Set(),
     envs,
+    newEnv,
   };
   let suites: Suite[] = [];
 
@@ -167,12 +169,12 @@ export function newTestTypeImpl(envs: Env<any>[]): any {
     suites[0]._testOptions = options;
   };
   test.extend = (env?: Env<any>) => {
-    const newTestType = newTestTypeImpl(env ? [...envs, env] : envs);
+    const newTestType = newTestTypeImpl(env ? [...envs, env] : envs, env);
     description.children.add(newTestType);
     return newTestType;
   };
   test.declare = () => {
-    const newTestType = newTestTypeImpl(envs);
+    const newTestType = newTestTypeImpl(envs, undefined);
     description.children.add(newTestType);
     return newTestType;
   };

diff --git a/src/workerRunner.ts b/src/workerRunner.ts
@@ -106,12 +106,16 @@ export class WorkerRunner extends EventEmitter {
   }
 
   private async _initEnvIfNeeded(envs: Env<any>[]) {
+    envs = [this._runList.env, ...envs];
+
     if (this._envRunner) {
-      // We rely on the fact that worker only receives tests with the same env list.
+      // We rely on the fact that worker only receives tests where
+      // environments with beforeAll/afterAll are the prefix of env list.
+      this._envRunner.update(envs);
       return;
     }
 
-    this._envRunner = new EnvRunner([this._runList.env, ...envs]);
+    this._envRunner = new EnvRunner(envs);
     // TODO: separate timeout for beforeAll?
     const result = await raceAgainstDeadline(this._envRunner.runBeforeAll(this._workerInfo, this._runList.options), this._deadline());
     this._workerArgs = result.result;
@@ -509,6 +513,10 @@ class EnvRunner {
     this.envs = [...envs];
   }
 
+  update(envs: Env[]) {
+    this.envs = [...envs];
+  }
+
   stop() {
     this._isStopped = true;
   }
@@ -532,6 +540,8 @@ class EnvRunner {
     const count = this.workerArgs.length;
     for (let index = count - 1; index >= 0; index--) {
       const args = this.workerArgs.pop();
+      if (this.envs.length <= index)
+        continue;
       const env = this.envs[index];
       if (env.afterAll) {
         try {
@@ -563,8 +573,10 @@ class EnvRunner {
     let error: Error | undefined;
     const count = this.testArgs.length;
     for (let index = count - 1; index >= 0; index--) {
-      const env = this.envs[index];
       const args = this.testArgs.pop();
+      if (this.envs.length <= index)
+        continue;
+      const env = this.envs[index];
       if (env.afterEach) {
         try {
           await env.afterEach(mergeObjects(testOptions, args), testInfo);

diff --git a/test/env.spec.ts b/test/env.spec.ts
@@ -225,3 +225,73 @@ test('should run sync env methods and hooks', async ({ runInlineTest }) => {
   });
   expect(passed).toBe(2);
 });
+
+test('should not create a new worker for environment with beforeEach only', async ({ runInlineTest }) => {
+  const result = await runInlineTest({
+    'a.test.ts': `
+      test('base test', async ({}, testInfo) => {
+        expect(testInfo.workerIndex).toBe(0);
+      });
+
+      const test2 = test.extend({
+        beforeEach() {
+          console.log('beforeEach-a');
+        }
+      });
+      test2('a test', async ({}, testInfo) => {
+        expect(testInfo.workerIndex).toBe(0);
+      });
+    `,
+    'b.test.ts': `
+      const test2 = test.extend({
+        beforeEach() {
+          console.log('beforeEach-b');
+        }
+      });
+      const test3 = test2.extend({
+        beforeEach() {
+          console.log('beforeEach-c');
+        }
+      });
+      test3('b test', async ({}, testInfo) => {
+        expect(testInfo.workerIndex).toBe(0);
+      });
+    `,
+  }, { workers: 1 });
+  expect(result.output).toContain('beforeEach-a');
+  expect(result.output).toContain('beforeEach-b');
+  expect(result.output).toContain('beforeEach-c');
+  expect(result.passed).toBe(3);
+});
+
+test('should create a new worker for environment with afterAll', async ({ runInlineTest }) => {
+  const result = await runInlineTest({
+    'a.test.ts': `
+      test('base test', async ({}, testInfo) => {
+        expect(testInfo.workerIndex).toBe(0);
+      });
+
+      const test2 = test.extend({
+        beforeAll() {
+          console.log('beforeAll-a');
+        }
+      });
+      test2('a test', async ({}, testInfo) => {
+        expect(testInfo.workerIndex).toBe(1);
+      });
+    `,
+    'b.test.ts': `
+      const test2 = test.extend({
+        beforeEach() {
+          console.log('beforeEach-b');
+        }
+      });
+      test2('b test', async ({}, testInfo) => {
+        expect(testInfo.workerIndex).toBe(0);
+      });
+    `,
+  }, { workers: 1 });
+  expect(result.output).toContain('beforeAll-a');
+  expect(result.output).toContain('beforeEach-b');
+  expect(result.passed).toBe(3);
+});