From 53208719c486f7c7f45982a609002b738e9fcd95 Mon Sep 17 00:00:00 2001
From: Nevenq <38163651+Nevenq@users.noreply.github.com>
Date: Thu, 23 Sep 2021 08:40:06 +0500
Subject: [PATCH] fix: Add browserPerWorker setting (#420)

* added browser per jest worker

* eslint error fixes

* removed unused variable

* Update packages/jest-environment-puppeteer/src/PuppeteerEnvironment.js

Co-authored-by: Tony Brix <tony@brix.ninja>

* added try catch for JSON.parse

* added doc for browser per worker

* review fixes

Co-authored-by: p.sem <p.sem@skbkontur.ru>
Co-authored-by: Tony Brix <tony@brix.ninja>
---
 README.md                                     |  1 +
 .../src/PuppeteerEnvironment.js               | 17 +++++++-
 .../jest-environment-puppeteer/src/global.js  | 41 +++++++++++++------
 3 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index a8c45906..c7f1581d 100644
--- a/README.md
+++ b/README.md
@@ -417,6 +417,7 @@ You can specify a `jest-puppeteer.config.js` at the root of the project or defin
 - `launch` <[object]> [All Puppeteer launch options](https://github.com/GoogleChrome/puppeteer/blob/master/docs/api.md#puppeteerlaunchoptions) can be specified in config. Since it is JavaScript, you can use all stuff you need, including environment.
 - `connect` <[object]> [All Puppeteer connect options](https://github.com/GoogleChrome/puppeteer/blob/master/docs/api.md#puppeteerconnectoptions) can be specified in config. This is an alternative to `launch` config, allowing you to connect to an already running instance of Chrome.
 - `server` <[Object]> Server options allowed by [jest-dev-server](https://github.com/smooth-code/jest-puppeteer/tree/master/packages/jest-dev-server)
+- `browserPerWorker` <[Boolean]> Allows to run tests for each [jest worker](https://jestjs.io/docs/cli#--maxworkersnumstring) in an individual browser.
 
 #### Example 1
 
diff --git a/packages/jest-environment-puppeteer/src/PuppeteerEnvironment.js b/packages/jest-environment-puppeteer/src/PuppeteerEnvironment.js
index 7ca85859..c1d1f1cf 100644
--- a/packages/jest-environment-puppeteer/src/PuppeteerEnvironment.js
+++ b/packages/jest-environment-puppeteer/src/PuppeteerEnvironment.js
@@ -12,6 +12,11 @@ const KEYS = {
   CONTROL_D: '\u0004',
   ENTER: '\r',
 }
+// JEST_WORKER_ID starts at 1
+const getWorkerIndex = () => process.env.JEST_WORKER_ID - 1
+
+const getEndpointIndex = () =>
+  Math.min(+process.env.BROWSERS_COUNT - 1, getWorkerIndex())
 
 class PuppeteerEnvironment extends NodeEnvironment {
   // Jest is not available here, so we have to reverse engineer
@@ -29,8 +34,16 @@ class PuppeteerEnvironment extends NodeEnvironment {
     const config = await readConfig()
     const puppeteer = getPuppeteer()
     this.global.puppeteerConfig = config
-
-    const wsEndpoint = process.env.PUPPETEER_WS_ENDPOINT
+    let wsEndpoint
+    try {
+      wsEndpoint = JSON.parse(process.env.PUPPETEER_WS_ENDPOINTS)[
+        getEndpointIndex()
+      ]
+    } catch (e) {
+      throw new Error(
+        `wsEndpoints parse error: ${e.message} in ${process.env.PUPPETEER_WS_ENDPOINTS}`,
+      )
+    }
     if (!wsEndpoint) {
       throw new Error('wsEndpoint not found')
     }
diff --git a/packages/jest-environment-puppeteer/src/global.js b/packages/jest-environment-puppeteer/src/global.js
index 42f3ec39..c668b08d 100644
--- a/packages/jest-environment-puppeteer/src/global.js
+++ b/packages/jest-environment-puppeteer/src/global.js
@@ -8,19 +8,33 @@ import {
 import chalk from 'chalk'
 import { readConfig, getPuppeteer } from './readConfig'
 
-let browser
+let browsers = []
 
 let didAlreadyRunInWatchMode = false
 
+async function openBrowser(puppeteer, config) {
+  if (config.connect) {
+    return puppeteer.connect(config.connect)
+  }
+  return puppeteer.launch(config.launch)
+}
+
 export async function setup(jestConfig = {}) {
   const config = await readConfig()
   const puppeteer = getPuppeteer()
-  if (config.connect) {
-    browser = await puppeteer.connect(config.connect)
-  } else {
-    browser = await puppeteer.launch(config.launch)
-  }
-  process.env.PUPPETEER_WS_ENDPOINT = browser.wsEndpoint()
+  const browsersCount =
+    config.browserPerWorker && !config.connect ? jestConfig.maxWorkers : 1
+  process.env.BROWSERS_COUNT = browsersCount
+
+  browsers = await Promise.all(
+    Array.from({ length: browsersCount }).map(() =>
+      openBrowser(puppeteer, config),
+    ),
+  )
+
+  const wsEndpoints = browsers.map((browser) => browser.wsEndpoint())
+
+  process.env.PUPPETEER_WS_ENDPOINTS = JSON.stringify(wsEndpoints)
 
   // If we are in watch mode, - only setupServer() once.
   if (jestConfig.watch || jestConfig.watchAll) {
@@ -60,11 +74,14 @@ export async function setup(jestConfig = {}) {
 export async function teardown(jestConfig = {}) {
   const config = await readConfig()
 
-  if (config.connect) {
-    await browser.disconnect()
-  } else {
-    await browser.close()
-  }
+  await Promise.all(
+    browsers.map((browser) => {
+      if (config.connect) {
+        return browser.disconnect()
+      }
+      return browser.close()
+    }),
+  )
 
   if (!jestConfig.watch && !jestConfig.watchAll) {
     await teardownServer()