From d846c056197b6f8e94dd90c7992119468f4b5e4b Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Thu, 5 Aug 2021 15:00:00 -0700
Subject: [PATCH] fix(test runner): make obtainWorker() resolve with null when
 stopping (#8018)

This ensures that we properly exit from `Dispatcher.run()`, print
epilogue and set the right exit code.
---
 src/test/dispatcher.ts                     | 10 +++++--
 tests/playwright-test/max-failures.spec.ts | 33 +++++++++++++++++++++-
 2 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/src/test/dispatcher.ts b/src/test/dispatcher.ts
index 87a91aa5e5..35b4d6fe08 100644
--- a/src/test/dispatcher.ts
+++ b/src/test/dispatcher.ts
@@ -71,11 +71,11 @@ export class Dispatcher {
       const testGroup = this._queue.shift()!;
       const requiredHash = testGroup.workerHash;
       let worker = await this._obtainWorker(testGroup);
-      while (!this._isStopped && worker.hash && worker.hash !== requiredHash) {
+      while (worker && worker.hash && worker.hash !== requiredHash) {
         worker.stop();
         worker = await this._obtainWorker(testGroup);
       }
-      if (this._isStopped)
+      if (this._isStopped || !worker)
         break;
       jobs.push(this._runJob(worker, testGroup));
     }
@@ -182,6 +182,8 @@ export class Dispatcher {
 
   async _obtainWorker(testGroup: TestGroup) {
     const claimWorker = (): Promise<Worker> | null => {
+      if (this._isStopped)
+        return null;
       // Use available worker.
       if (this._freeWorkers.length)
         return Promise.resolve(this._freeWorkers.pop()!);
@@ -199,7 +201,7 @@ export class Dispatcher {
       await new Promise<void>(f => this._workerClaimers.push(f));
       worker = claimWorker();
     }
-    return worker!;
+    return worker;
   }
 
   async _notifyWorkerClaimer() {
@@ -294,6 +296,8 @@ export class Dispatcher {
         worker.stop();
       await result;
     }
+    while (this._workerClaimers.length)
+      this._workerClaimers.shift()!();
   }
 
   private _hasReachedMaxFailures() {
diff --git a/tests/playwright-test/max-failures.spec.ts b/tests/playwright-test/max-failures.spec.ts
index 2a931d832a..6a456032d3 100644
--- a/tests/playwright-test/max-failures.spec.ts
+++ b/tests/playwright-test/max-failures.spec.ts
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { test, expect } from './playwright-test-fixtures';
+import { test, expect, stripAscii } from './playwright-test-fixtures';
 
 test('max-failures should work', async ({ runInlineTest }) => {
   const result = await runInlineTest({
@@ -112,3 +112,34 @@ test('max-failures should stop workers', async ({ runInlineTest }) => {
   expect(result.output).toContain('%%interrupted');
   expect(result.output).not.toContain('%%skipped');
 });
+
+test('max-failures should properly shutdown', async ({ runInlineTest }) => {
+  const result = await runInlineTest({
+    'playwright.config.ts': `
+      const config = {
+        testDir: './',
+        maxFailures: 1,
+      };
+      export default config;
+    `,
+    'test1.spec.ts': `
+      const { test } = pwt;
+      test.describe('spec 1', () => {
+        test('test 1', async () => {
+          expect(false).toBeTruthy()
+        })
+      });
+    `,
+    'test2.spec.ts': `
+      const { test } = pwt;
+      test.describe('spec 2', () => {
+        test('test 2', () => {
+          expect(true).toBeTruthy()
+        })
+      });
+    `,
+  }, { workers: 1 });
+  expect(result.exitCode).toBe(1);
+  expect(result.failed).toBe(1);
+  expect(stripAscii(result.output)).toContain('expect(false).toBeTruthy()');
+});