From 5f4feacdb23acd61a93f5ec8ab3f0abaf40c8088 Mon Sep 17 00:00:00 2001
From: Nikita Skovoroda <chalkerx@gmail.com>
Date: Mon, 15 Jul 2024 19:53:41 +0300
Subject: [PATCH] feat: test.concurrent (only block-level)

---
 __test__/jest/concurrent.js | 35 +++++++++++++++++++++++++++++++++++
 src/jest.js                 | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 __test__/jest/concurrent.js

diff --git a/__test__/jest/concurrent.js b/__test__/jest/concurrent.js
new file mode 100644
index 0000000..e1c703c
--- /dev/null
+++ b/__test__/jest/concurrent.js
@@ -0,0 +1,35 @@
+const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms))
+
+describe('concurrent', () => {
+  let i = 0
+  test.concurrent(async () => {
+    i++
+    await sleep(50)
+    expect(i).toBe(2)
+    await sleep(50)
+    i--
+  })
+  test.concurrent(async () => {
+    i++
+    await sleep(50)
+    expect(i).toBe(2)
+    await sleep(50)
+    i--
+  })
+
+  let j = 0
+  test.concurrent(async () => {
+    j++
+    await sleep(50)
+    expect(j).toBe(1)
+    await sleep(50)
+    j--
+  })
+  test.concurrent(async () => {
+    j++
+    await sleep(50)
+    expect(j).toBe(1)
+    await sleep(50)
+    j--
+  })
+})
diff --git a/src/jest.js b/src/jest.js
index dc735dd..dcd9ae5 100644
--- a/src/jest.js
+++ b/src/jest.js
@@ -15,6 +15,7 @@ const { getCallerLocation, installLocationInNextTest } = createCallerLocationHoo
 expect.extend(matchers)
 
 let defaultTimeout = jestConfig().testTimeout // overridable via jest.setTimeout()
+const defaultConcurrency = jestConfig().maxConcurrency
 
 function parseArgs(list, targs) {
   if (!(Object.isFrozen(list) && list.length === targs.length + 1)) return list // template check
@@ -79,7 +80,33 @@ const makeEach =
 
 const forceExit = process.execArgv.map((x) => x.replaceAll('_', '-')).includes('--test-force-exit')
 
-const describe = (...args) => nodeDescribe(...args)
+const inConcurrent = []
+const concurrent = []
+const describe = (...args) => {
+  const fn = args.pop()
+  const optionsConcurrent = args?.at(-1)?.concurrency > 1
+  if (optionsConcurrent) inConcurrent.push(fn)
+  const result = nodeDescribe(...args, async () => {
+    const res = fn()
+
+    // We do only block-level concurrency, not file-level
+    if (concurrent.length === 1) {
+      test(...concurrent[0])
+      concurrent.length = 0
+    } else if (concurrent.length > 0) {
+      const queue = [...concurrent]
+      concurrent.length = 0
+      describe('concurrent', { concurrency: defaultConcurrency }, () => {
+        for (const args of queue) test(...args)
+      })
+    }
+
+    return res
+  })
+  if (optionsConcurrent) inConcurrent.pop()
+  return result
+}
+
 const test = (name, fn, testTimeout) => {
   const timeout = testTimeout ?? defaultTimeout
   installLocationInNextTest(getCallerLocation())
@@ -99,6 +126,8 @@ Also, using expect.assertions() to ensure the planned number of assertions is be
 
 describe.each = makeEach(describe)
 test.each = makeEach(test)
+test.concurrent = (...args) => (inConcurrent.length > 0 ? test(...args) : concurrent.push(args))
+test.concurrent.each = makeEach(test.concurrent)
 describe.skip = (...args) => nodeDescribe.skip(...args)
 test.skip = (...args) => nodeTest.skip(...args)