feat(ci): Balance frontend tests

jest.config.ts

@@ -8,8 +8,15 @@ import type {Config} from '@jest/types';
 
 import babelConfig from './babel.config';
 
-const {CI, JEST_TESTS, CI_NODE_TOTAL, CI_NODE_INDEX, GITHUB_PR_SHA, GITHUB_PR_REF} =
-  process.env;
+const {
+  CI,
+  JEST_TESTS,
+  JEST_TEST_BALANCER,
+  CI_NODE_TOTAL,
+  CI_NODE_INDEX,
+  GITHUB_PR_SHA,
+  GITHUB_PR_REF,
+} = process.env;
 
 /**
  * In CI we may need to shard our jest tests so that we can parellize the test runs
@@ -20,25 +27,111 @@ const {CI, JEST_TESTS, CI_NODE_TOTAL, CI_NODE_INDEX, GITHUB_PR_SHA, GITHUB_PR_RE
  */
 let testMatch: string[] | undefined;
 
+const BALANCE_RESULTS_PATH = path.resolve(
+  __dirname,
+  'tests',
+  'js',
+  'test-balancer',
+  'jest-balance.json'
+);
+
+/**
+ * Given a Map of <testName, testRunTime> and a number of total groups, split the
+ * tests into n groups whose total test run times should be roughly equal
+ *
+ * The source results should be sorted with the slowest tests first. We insert
+ * the test into the smallest group on each interation. This isn't perfect, but
+ * should be good enough.
+ *
+ * Returns a map of <testName, groupIndex>
+ */
+function balancer(
+  allTests: string[],
+  source: Record<string, number>,
+  numberGroups: number
+) {
+  const results = new Map<string, number>();
+  const totalRunTimes = Array(numberGroups).fill(0);
+
+  /**
+   * Find the index of the smallest group (totalRunTimes)
+   */
+  function findSmallestGroup() {
+    let index = 0;
+    let smallestRunTime = null;
+    for (let i = 0; i < totalRunTimes.length; i++) {
+      const runTime = totalRunTimes[i];
+
+      if (!smallestRunTime || runTime <= smallestRunTime) {
+        smallestRunTime = totalRunTimes[i];
+        index = i;
+      }
+
+      if (runTime === 0) {
+        break;
+      }
+    }
+
+    return index;
+  }
+
+  /**
+   * We may not have a duration for all tests (e.g. a test that was just added)
+   * as the `source` needs to be generated
+   */
+  for (const test of allTests) {
+    const index = findSmallestGroup();
+    results.set(test, index);
+
+    if (source[test] !== undefined) {
+      totalRunTimes[index] = totalRunTimes[index] + source[test];
+    }
+  }
+
+  return results;
+}
+
 if (
   JEST_TESTS &&
   typeof CI_NODE_TOTAL !== 'undefined' &&
   typeof CI_NODE_INDEX !== 'undefined'
 ) {
+  let balance: null | Record<string, number> = null;
+
+  try {
+    balance = require(BALANCE_RESULTS_PATH);
+  } catch (err) {
+    // Just ignore if balance results doesn't exist
+  }
+
   // Taken from https://github.com/facebook/jest/issues/6270#issue-326653779
-  const envTestList = JSON.parse(JEST_TESTS) as string[];
+  const envTestList = JSON.parse(JEST_TESTS).map(file =>
+    file.replace(__dirname, '')
+  ) as string[];
   const tests = envTestList.sort((a, b) => b.localeCompare(a));
 
   const nodeTotal = Number(CI_NODE_TOTAL);
   const nodeIndex = Number(CI_NODE_INDEX);
 
-  const length = tests.length;
-  const size = Math.floor(length / nodeTotal);
-  const remainder = length % nodeTotal;
-  const offset = Math.min(nodeIndex, remainder) + nodeIndex * size;
-  const chunk = size + (nodeIndex < remainder ? 1 : 0);
+  if (balance) {
+    const results = balancer(envTestList, balance, nodeTotal);
 
-  testMatch = tests.slice(offset, offset + chunk);
+    testMatch = [
+      // First, we only want the tests that we have test durations for and belong
+      // to the current node's index
+      ...Object.entries(Object.fromEntries(results))
+        .filter(([, index]) => index === nodeIndex)
+        .map(([test]) => `${path.join(__dirname, test)}`),
+    ];
+  } else {
+    const length = tests.length;
+    const size = Math.floor(length / nodeTotal);
+    const remainder = length % nodeTotal;
+    const offset = Math.min(nodeIndex, remainder) + nodeIndex * size;
+    const chunk = size + (nodeIndex < remainder ? 1 : 0);
+
+    testMatch = tests.slice(offset, offset + chunk);
+  }
 }
 
 const config: Config.InitialOptions = {
@@ -93,6 +186,13 @@ const config: Config.InitialOptions = {
         outputName: 'jest.junit.xml',
       },
     ],
+    [
+      '<rootDir>/tests/js/test-balancer',
+      {
+        enabled: !!JEST_TEST_BALANCER,
+        resultsPath: BALANCE_RESULTS_PATH,
+      },
+    ],
   ],
 
   testRunner: 'jest-circus/runner',

tests/js/spec/components/checkbox.spec.tsx

@@ -1,11 +1,12 @@
-import {mountWithTheme} from 'sentry-test/reactTestingLibrary';
+import {mountWithTheme, screen} from 'sentry-test/reactTestingLibrary';
 
 import Checkbox from 'sentry/components/checkbox';
 
 describe('Checkbox', function () {
-  it('renders', function () {
+  it('renders', async function () {
     const {container} = mountWithTheme(<Checkbox onChange={() => {}} />);
 
+    expect(await screen.findByRole('checkbox')).toBeInTheDocument();
     expect(container).toSnapshot();
   });
 });

tests/js/spec/components/circleIndicator.spec.tsx

@@ -1,10 +1,14 @@
-import {mountWithTheme} from 'sentry-test/reactTestingLibrary';
+import {mountWithTheme, screen} from 'sentry-test/reactTestingLibrary';
 
 import CircleIndicator from 'sentry/components/circleIndicator';
 
 describe('CircleIndicator', function () {
-  it('renders', function () {
-    const {container} = mountWithTheme(<CircleIndicator />);
+  it('renders', async function () {
+    const {container} = mountWithTheme(
+      <CircleIndicator data-test-id="circleIndicator" />
+    );
+
+    expect(await screen.findByTestId('circleIndicator')).toBeInTheDocument();
     expect(container).toSnapshot();
   });
 });

tests/js/spec/components/returnButton.spec.tsx

@@ -1,10 +1,11 @@
-import {mountWithTheme} from 'sentry-test/reactTestingLibrary';
+import {mountWithTheme, screen} from 'sentry-test/reactTestingLibrary';
 
 import ReturnButton from 'sentry/views/settings/components/forms/returnButton';
 
 describe('returnButton', function () {
-  it('renders', function () {
-    const {container} = mountWithTheme(<ReturnButton />);
+  it('renders', async function () {
+    const {container} = mountWithTheme(<ReturnButton data-test-id="returnButton" />);
+    expect(await screen.findByTestId('returnButton')).toBeInTheDocument();
     expect(container).toSnapshot();
   });
 });

tests/js/spec/views/admin/adminBuffer.spec.jsx

@@ -12,6 +12,8 @@ describe('AdminBuffer', function () {
           router: TestStubs.router(),
         },
       });
+
+      expect(wrapper.find('LoadingIndicator')).toHaveLength(2);
       expect(wrapper).toSnapshot();
     });
   });

tests/js/test-balancer/index.js

@@ -0,0 +1,34 @@
+/* eslint-env node */
+/* eslint import/no-nodejs-modules:0 */
+const fs = require('fs/promises');
+
+class TestBalancer {
+  constructor(globalConfig, options) {
+    this._globalConfig = globalConfig;
+    this._options = options;
+
+    this.results = new Map();
+  }
+
+  onTestFileResult(test) {
+    const path = test.path.replace(this._globalConfig.rootDir, '');
+    this.results.set(path, test.duration);
+  }
+
+  onRunComplete(_contexts, results) {
+    // results.success always returns false for me?
+    if (
+      results.numTotalTests === 0 ||
+      results.numFailedTests > 0 ||
+      !this._options.enabled ||
+      !this._options.resultsPath
+    ) {
+      return;
+    }
+
+    const data = JSON.stringify(Object.fromEntries(this.results), null, 2);
+    fs.writeFile(this._options.resultsPath, data);
+  }
+}
+
+module.exports = TestBalancer;