From 1fe84a038bad3f82710734c1941608b12a60e2bc Mon Sep 17 00:00:00 2001
From: Mario Nebl <mnebl@atlassian.com>
Date: Wed, 19 Oct 2022 22:01:41 +1100
Subject: [PATCH] fix: distribute rest files over shards #13027

---
 package.json                                  |  5 +++-
 .../src/__tests__/test_sequencer.test.ts      | 19 ++++++++++++-
 packages/jest-test-sequencer/src/index.ts     | 27 ++++++++++++++++---
 yarn.lock                                     | 21 +++++++++++++++
 4 files changed, 67 insertions(+), 5 deletions(-)

diff --git a/package.json b/package.json
index be2d5a9654d7..7164c3cb0f67 100644
--- a/package.json
+++ b/package.json
@@ -180,5 +180,8 @@
     "jest": "workspace:^",
     "jest-environment-node": "workspace:^"
   },
-  "packageManager": "yarn@3.2.4"
+  "packageManager": "yarn@3.2.4",
+  "dependencies": {
+    "terminal-link": "^3.0.0"
+  }
 }
diff --git a/packages/jest-test-sequencer/src/__tests__/test_sequencer.test.ts b/packages/jest-test-sequencer/src/__tests__/test_sequencer.test.ts
index 88b5bb4b9c17..87a6b8157f91 100644
--- a/packages/jest-test-sequencer/src/__tests__/test_sequencer.test.ts
+++ b/packages/jest-test-sequencer/src/__tests__/test_sequencer.test.ts
@@ -368,6 +368,23 @@ test('returns expected 100/8 shards', async () => {
   );
 
   expect(shards.map(shard => shard.length)).toEqual([
-    13, 13, 13, 13, 13, 13, 13, 9,
+    13, 13, 13, 13, 12, 12, 12, 12,
+  ]);
+});
+
+test('returns expected 55/12 shards', async () => {
+  const allTests = toTests(new Array(55).fill(true).map((_, i) => `/${i}.js`));
+
+  const shards = await Promise.all(
+    new Array(12).fill(true).map((_, i) =>
+      sequencer.shard(allTests, {
+        shardCount: 12,
+        shardIndex: i + 1,
+      }),
+    ),
+  );
+
+  expect(shards.map(shard => shard.length)).toEqual([
+    5, 5, 5, 5, 5, 5, 5, 4, 4, 4, 4, 4,
   ]);
 });
diff --git a/packages/jest-test-sequencer/src/index.ts b/packages/jest-test-sequencer/src/index.ts
index f55b5ea4c1f2..a0c6161d5666 100644
--- a/packages/jest-test-sequencer/src/index.ts
+++ b/packages/jest-test-sequencer/src/index.ts
@@ -72,6 +72,19 @@ export default class TestSequencer {
     return cache;
   }
 
+  _shardPosition(options: ShardOptions & {suiteLength: number}): number {
+    const shardRest = options.suiteLength % options.shardCount;
+    const ratio = options.suiteLength / options.shardCount;
+
+    return new Array(options.shardIndex)
+      .fill(true)
+      .reduce((acc, _, shardIndex) => {
+        const dangles = shardIndex < shardRest;
+        const shardSize = dangles ? Math.ceil(ratio) : Math.floor(ratio);
+        return acc + shardSize;
+      }, 0);
+  }
+
   /**
    * Select tests for shard requested via --shard=shardIndex/shardCount
    * Sharding is applied before sorting
@@ -97,9 +110,17 @@ export default class TestSequencer {
     tests: Array<Test>,
     options: ShardOptions,
   ): Array<Test> | Promise<Array<Test>> {
-    const shardSize = Math.ceil(tests.length / options.shardCount);
-    const shardStart = shardSize * (options.shardIndex - 1);
-    const shardEnd = shardSize * options.shardIndex;
+    const shardStart = this._shardPosition({
+      shardCount: options.shardCount,
+      shardIndex: options.shardIndex - 1,
+      suiteLength: tests.length,
+    });
+
+    const shardEnd = this._shardPosition({
+      shardCount: options.shardCount,
+      shardIndex: options.shardIndex,
+      suiteLength: tests.length,
+    });
 
     return tests
       .map(test => {
diff --git a/yarn.lock b/yarn.lock
index 7c9a5d2160d9..932957b87719 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -2828,6 +2828,7 @@ __metadata:
     strip-ansi: ^6.0.0
     strip-json-comments: ^3.1.1
     tempy: ^1.0.0
+    terminal-link: ^3.0.0
     ts-node: ^10.5.0
     typescript: ^4.8.2
     which: ^2.0.1
@@ -19350,6 +19351,16 @@ __metadata:
   languageName: node
   linkType: hard
 
+"supports-hyperlinks@npm:^2.2.0":
+  version: 2.3.0
+  resolution: "supports-hyperlinks@npm:2.3.0"
+  dependencies:
+    has-flag: ^4.0.0
+    supports-color: ^7.0.0
+  checksum: 9ee0de3c8ce919d453511b2b1588a8205bd429d98af94a01df87411391010fe22ca463f268c84b2ce2abad019dfff8452aa02806eeb5c905a8d7ad5c4f4c52b8
+  languageName: node
+  linkType: hard
+
 "supports-preserve-symlinks-flag@npm:^1.0.0":
   version: 1.0.0
   resolution: "supports-preserve-symlinks-flag@npm:1.0.0"
@@ -19499,6 +19510,16 @@ __metadata:
   languageName: node
   linkType: hard
 
+"terminal-link@npm:^3.0.0":
+  version: 3.0.0
+  resolution: "terminal-link@npm:3.0.0"
+  dependencies:
+    ansi-escapes: ^5.0.0
+    supports-hyperlinks: ^2.2.0
+  checksum: 85a78ae50a2cd3c43df25922e7572f1008c92b1ea98c6c4579bbbe02fa54677a487123c3cae44fecd1a36cac782d0be2cec212a916818abb2b4df6fbb8eed341
+  languageName: node
+  linkType: hard
+
 "terser-webpack-plugin@npm:^5.1.3, terser-webpack-plugin@npm:^5.3.3":
   version: 5.3.6
   resolution: "terser-webpack-plugin@npm:5.3.6"