vercel/utils/chunk-tests.js

const child_process = require('child_process');
const path = require('path');

const NUMBER_OF_CHUNKS = 5;
const MINIMUM_PER_CHUNK = 1;
const runnersMap = new Map([['test-integration-once', ['ubuntu-latest']]]);

async function getChunkedTests() {
  const scripts = ['test-integration-once'];
  const rootPath = path.resolve(__dirname, '..');

  const dryRunText = (
    await turbo([`run`, ...scripts, `--cache-dir=.turbo`, '--', '--listTests'])
  ).toString('utf8');

  /**
   * @typedef {string} TestPath
   * @type {{ [package: string]: { [script: string]: TestPath[] } }}
   */
  const testsToRun = {};

  dryRunText
    .split('\n')
    .flatMap(line => {
      const [packageAndScriptName, possiblyPath] = line.split(' ');
      const [packageName, scriptName] = packageAndScriptName.split(':');

      if (!packageName || !scriptName || !possiblyPath?.includes(rootPath)) {
        return [];
      }

      return { testPath: possiblyPath, scriptName, packageName };
    })
    .forEach(({ testPath, scriptName, packageName }) => {
      const [, packageDir, shortPackageName] = testPath
        .replace(rootPath, '')
        .split(path.sep);

      const packagePath =
        path.join(packageDir, shortPackageName) + ',' + packageName;
      testsToRun[packagePath] = testsToRun[packagePath] || {};
      testsToRun[packagePath][scriptName] =
        testsToRun[packagePath][scriptName] || [];
      testsToRun[packagePath][scriptName].push(testPath);
    });

  const chunkedTests = Object.entries(testsToRun).flatMap(
    ([packagePathAndName, scriptNames]) => {
      const [packagePath, packageName] = packagePathAndName.split(',');
      return Object.entries(scriptNames).flatMap(([scriptName, testPaths]) => {
        return intoChunks(NUMBER_OF_CHUNKS, testPaths).flatMap(
          (chunk, chunkNumber, allChunks) => {
            const runners = runnersMap.get(scriptName) || ['ubuntu-latest'];

            return runners.map(runner => {
              return {
                runner,
                packagePath,
                packageName,
                scriptName,
                testPaths: chunk.map(testFile =>
                  path.relative(
                    path.join(__dirname, '../', packagePath),
                    testFile
                  )
                ),
                chunkNumber: chunkNumber + 1,
                allChunksLength: allChunks.length,
              };
            });
          }
        );
      });
    }
  );

  return chunkedTests;
}

async function turbo(args) {
  const chunks = [];
  try {
    await new Promise((resolve, reject) => {
      const spawned = child_process.spawn(`yarn`, ['turbo', '--', ...args], {
        cwd: path.resolve(__dirname, '..'),
        env: {
          ...process.env,
          YARN_SILENT: '1',
        },
      });
      spawned.stdout.on('data', data => {
        chunks.push(data);
        process.stderr.write(data);
      });
      spawned.stderr.on('data', data => {
        process.stderr.write(data);
      });
      spawned.on('close', code => {
        if (code !== 0) {
          reject(new Error(`Turbo exited with code ${code}`));
        } else {
          resolve();
        }
      });
    });
    return Buffer.concat(chunks);
  } catch (e) {
    console.error(e.message);
    process.exit(1);
  }
}

/**
 * @template T
 * @param {number} totalChunks maximum number of chunks
 * @param {T[]} values
 * @returns {T[][]}
 */
function intoChunks(totalChunks, arr) {
  const chunkSize = Math.max(
    MINIMUM_PER_CHUNK,
    Math.ceil(arr.length / totalChunks)
  );
  const chunks = [];
  for (let i = 0; i < totalChunks; i++) {
    chunks.push(arr.slice(i * chunkSize, (i + 1) * chunkSize));
  }
  return chunks.filter(x => x.length > 0);
}

async function main() {
  try {
    const chunks = await getChunkedTests(process.argv[2]);
    // TODO: pack and build the runtimes for each package and cache it so we only deploy it once
    console.log(JSON.stringify(chunks));
  } catch (e) {
    console.error(e);
    process.exitCode = 1;
  }
}

if (module === require.main || !module.parent) {
  main();
}

module.exports = {
  intoChunks,
  NUMBER_OF_CHUNKS,
};