Shard serialized server CI suites in PR workflow

Co-Authored-By: Paperclip <noreply@paperclip.ing>

scripts/run-vitest-stable.mjs

@@ -45,6 +45,9 @@ const additionalSerializedServerTests = new Set([
   "server/src/__tests__/routines-e2e.test.ts",
 ]);
 let invocationIndex = 0;
+const serializedModeName = "serialized";
+const generalModeName = "general";
+const allModeName = "all";
 
 function walk(dir) {
   const entries = readdirSync(dir);
@@ -77,6 +80,130 @@ function isRouteOrAuthzTest(file) {
   return additionalSerializedServerTests.has(file);
 }
 
+function fail(message) {
+  console.error(`[test:run] ${message}`);
+  process.exit(1);
+}
+
+function readOptionValue(argv, index, argName) {
+  const value = argv[index + 1];
+  if (!value) {
+    fail(`Missing value for ${argName}`);
+  }
+
+  return value;
+}
+
+function parseNonNegativeInteger(value, argName) {
+  const parsed = Number.parseInt(value, 10);
+  if (!Number.isInteger(parsed) || parsed < 0) {
+    fail(`${argName} must be a non-negative integer. Received "${value}".`);
+  }
+
+  return parsed;
+}
+
+function parsePositiveInteger(value, argName) {
+  const parsed = Number.parseInt(value, 10);
+  if (!Number.isInteger(parsed) || parsed < 1) {
+    fail(`${argName} must be a positive integer. Received "${value}".`);
+  }
+
+  return parsed;
+}
+
+function parseCliOptions(argv) {
+  let mode = allModeName;
+  let shardIndex = null;
+  let shardCount = null;
+  let dryRun = false;
+
+  for (let index = 0; index < argv.length; index += 1) {
+    const arg = argv[index];
+    if (arg === "--") {
+      continue;
+    }
+
+    if (arg === "--mode") {
+      mode = readOptionValue(argv, index, arg);
+      index += 1;
+      continue;
+    }
+
+    if (arg.startsWith("--mode=")) {
+      mode = arg.slice("--mode=".length);
+      continue;
+    }
+
+    if (arg === "--shard-index") {
+      shardIndex = parseNonNegativeInteger(readOptionValue(argv, index, arg), arg);
+      index += 1;
+      continue;
+    }
+
+    if (arg.startsWith("--shard-index=")) {
+      shardIndex = parseNonNegativeInteger(arg.slice("--shard-index=".length), "--shard-index");
+      continue;
+    }
+
+    if (arg === "--shard-count") {
+      shardCount = parsePositiveInteger(readOptionValue(argv, index, arg), arg);
+      index += 1;
+      continue;
+    }
+
+    if (arg.startsWith("--shard-count=")) {
+      shardCount = parsePositiveInteger(arg.slice("--shard-count=".length), "--shard-count");
+      continue;
+    }
+
+    if (arg === "--dry-run") {
+      dryRun = true;
+      continue;
+    }
+
+    fail(`Unknown argument "${arg}".`);
+  }
+
+  if (!new Set([allModeName, generalModeName, serializedModeName]).has(mode)) {
+    fail(`Unknown mode "${mode}". Expected one of: ${allModeName}, ${generalModeName}, ${serializedModeName}.`);
+  }
+
+  if ((shardIndex === null) !== (shardCount === null)) {
+    fail("--shard-index and --shard-count must be provided together.");
+  }
+
+  if (mode !== serializedModeName && shardIndex !== null) {
+    fail("--shard-index/--shard-count are only valid with --mode serialized.");
+  }
+
+  if (mode === serializedModeName) {
+    const resolvedShardCount = shardCount ?? 1;
+    const resolvedShardIndex = shardIndex ?? 0;
+    if (resolvedShardIndex >= resolvedShardCount) {
+      fail(`--shard-index must be less than --shard-count. Received ${resolvedShardIndex} of ${resolvedShardCount}.`);
+    }
+
+    return {
+      mode,
+      shardIndex: resolvedShardIndex,
+      shardCount: resolvedShardCount,
+      dryRun,
+    };
+  }
+
+  return {
+    mode,
+    shardIndex: null,
+    shardCount: null,
+    dryRun,
+  };
+}
+
+function selectSerializedSuites(routeTests, shardIndex, shardCount) {
+  return routeTests.filter((_, index) => index % shardCount === shardIndex);
+}
+
 function runVitest(args, label) {
   console.log(`\n[test:run] ${label}`);
   invocationIndex += 1;
@@ -103,6 +230,38 @@ function runVitest(args, label) {
   }
 }
 
+function runGeneralSuites(routeTests) {
+  const excludeRouteArgs = routeTests.flatMap((file) => ["--exclude", file.serverPath]);
+  for (const project of nonServerProjects) {
+    runVitest(["--project", project], `non-server project ${project}`);
+  }
+
+  runVitest(
+    ["--project", "@paperclipai/server", ...excludeRouteArgs],
+    `server suites excluding ${routeTests.length} serialized suites`,
+  );
+}
+
+function runSerializedSuites(routeTests, shardIndex, shardCount) {
+  const shardTests = selectSerializedSuites(routeTests, shardIndex, shardCount);
+  console.log(
+    `\n[test:run] serialized shard ${shardIndex + 1}/${shardCount} running ${shardTests.length} of ${routeTests.length} suites`,
+  );
+
+  for (const routeTest of shardTests) {
+    runVitest(
+      [
+        "--project",
+        "@paperclipai/server",
+        routeTest.repoPath,
+        "--pool=forks",
+        "--poolOptions.forks.isolate=true",
+      ],
+      routeTest.repoPath,
+    );
+  }
+}
+
 const routeTests = walk(serverTestsDir)
   .filter((file) => isRouteOrAuthzTest(toRepoPath(file)))
   .map((file) => ({
@@ -111,25 +270,32 @@ const routeTests = walk(serverTestsDir)
   }))
   .sort((a, b) => a.repoPath.localeCompare(b.repoPath));
 
-const excludeRouteArgs = routeTests.flatMap((file) => ["--exclude", file.serverPath]);
-for (const project of nonServerProjects) {
-  runVitest(["--project", project], `non-server project ${project}`);
-}
-
-runVitest(
-  ["--project", "@paperclipai/server", ...excludeRouteArgs],
-  `server suites excluding ${routeTests.length} serialized suites`,
-);
-
-for (const routeTest of routeTests) {
-  runVitest(
-    [
-      "--project",
-      "@paperclipai/server",
-      routeTest.repoPath,
-      "--pool=forks",
-      "--poolOptions.forks.isolate=true",
-    ],
-    routeTest.repoPath,
+const options = parseCliOptions(process.argv.slice(2));
+if (options.dryRun) {
+  const serializedSuites =
+    options.mode === serializedModeName
+      ? selectSerializedSuites(routeTests, options.shardIndex, options.shardCount)
+      : routeTests;
+  console.log(
+    JSON.stringify(
+      {
+        mode: options.mode,
+        shardIndex: options.shardIndex,
+        shardCount: options.shardCount,
+        serializedSuiteCount: routeTests.length,
+        selectedSerializedSuites: serializedSuites.map((routeTest) => routeTest.repoPath),
+      },
+      null,
+      2,
+    ),
   );
+  process.exit(0);
+}
+
+if (options.mode === generalModeName || options.mode === allModeName) {
+  runGeneralSuites(routeTests);
+}
+
+if (options.mode === serializedModeName || options.mode === allModeName) {
+  runSerializedSuites(routeTests, options.shardIndex ?? 0, options.shardCount ?? 1);
 }