test: enforce wheelhouse budgets for each test lane

Author: jdalton

.config/repo/socket-wheelhouse.json

@@ -174,6 +174,19 @@
     ]
   },
   "workspace": {
-    "catalogDriftIgnore": ["typescript"]
+    "catalogDriftIgnore": [
+      "typescript"
+    ]
+  },
+  "vitest": {
+    "lanes": {
+      "mid": [
+        "packages/cli/test/unit/commands/**",
+        "packages/cli/test/unit/meow.test.mts"
+      ],
+      "slow": [
+        "packages/cli/test/integration/**"
+      ]
+    }
   }
 }

packages/cli/package.json

@@ -64,20 +64,23 @@
     "e2e:js": "node scripts/e2e.mts --js",
     "e2e:sea": "node scripts/e2e.mts --sea",
     "e2e:all": "node scripts/e2e.mts --all",
-    "test": "run-s check test:prepare test:unit test:validate",
+    "test": "run-s check test:prepare test:unit:all test:validate",
     "test:prepare": "pnpm build && del-cli 'test/**/node_modules'",
     "test:fuzz": "node scripts/repo/fuzz.mts",
     "test:unit": "node --import=./scripts/load.mts scripts/test-wrapper.mts",
     "test:unit:update": "node --import=./scripts/load.mts scripts/test-wrapper.mts --update",
     "test:unit:coverage": "node --import=./scripts/load.mts scripts/test-wrapper.mts --coverage",
     "test:validate": "node --import=./scripts/load.mts scripts/validate-tests.mts",
-    "test-ci": "run-s test:prepare test:unit test:validate",
+    "test-ci": "run-s test:prepare test:unit:all test:validate",
     "test-pre-commit": "cross-env PRE_COMMIT=1 pnpm test",
     "update": "node ../../scripts/fleet/update.mts",
     "verify": "node scripts/verify-package.mts",
     "wasm": "node scripts/wasm.mts",
     "wasm:build": "node scripts/wasm.mts --build",
-    "wasm:download": "node scripts/wasm.mts --download"
+    "wasm:download": "node scripts/wasm.mts --download",
+    "test:unit:all": "node --import=./scripts/load.mts scripts/test-wrapper.mts --all",
+    "test:integration": "node --import=./scripts/load.mts scripts/test-wrapper.mts --config vitest.integration.config.mts --lane slow",
+    "test:mid": "node --import=./scripts/load.mts scripts/test-wrapper.mts --lane mid"
   },
   "dependencies": {
     "@socketsecurity/odai": "0.2.1",

packages/cli/scripts/test-lanes.mts

@@ -0,0 +1,79 @@
+import { readFileSync } from 'node:fs'
+import path from 'node:path'
+
+import { normalizePath } from '@socketsecurity/lib-stable/paths/normalize'
+
+import {
+  configuredFastBudgetMs,
+  coverBudgetMs,
+  laneBudgetMs,
+} from '../../../scripts/fleet/constants/test-budget.mts'
+import type { TestLane } from '../../../scripts/fleet/constants/test-budget.mts'
+import { extractLane } from '../../../scripts/fleet/test-runner/cli-args.mts'
+
+export function resolvePackageTestScope(
+  args: readonly string[],
+  repoRoot: string,
+) {
+  const parsed = extractLane(args)
+  const all = parsed.rest.includes('--all')
+  const rest = parsed.rest.filter(arg => arg !== '--all')
+  const requested =
+    parsed.lane ?? (!all && rest.length === 0 ? 'fast' : undefined)
+  const lane =
+    requested === 'fast' || requested === 'mid' || requested === 'slow'
+      ? requested
+      : undefined
+  return {
+    __proto__: null,
+    args: rest,
+    lane,
+    timeout: lane
+      ? laneBudgetMs(lane, { configuredFast: configuredFastBudgetMs(repoRoot) })
+      : coverBudgetMs(),
+  }
+}
+
+export function packageTestGlobs(
+  globs: readonly string[],
+  packagePath: string,
+): string[] {
+  const prefix = `${normalizePath(packagePath).replace(/\/$/u, '')}/`
+  return globs
+    .map(glob => normalizePath(glob))
+    .filter(glob => glob.startsWith(prefix))
+    .map(glob => glob.slice(prefix.length))
+}
+
+export function readPackageTestLanes(repoRoot: string, packagePath: string) {
+  const settings: {
+    vitest?:
+      | { lanes?: Partial<Record<TestLane, string[]>> | undefined }
+      | undefined
+  } = JSON.parse(
+    readFileSync(
+      path.join(repoRoot, '.config/repo/socket-wheelhouse.json'),
+      'utf8',
+    ),
+  )
+  return {
+    __proto__: null,
+    mid: packageTestGlobs(settings.vitest?.lanes?.mid ?? [], packagePath),
+    slow: packageTestGlobs(settings.vitest?.lanes?.slow ?? [], packagePath),
+  }
+}
+
+export function selectPackageTestGlobs(
+  lane: string | undefined,
+  lanes: { mid: string[]; slow: string[] },
+) {
+  const selected =
+    lane === 'mid' ? lanes.mid : lane === 'slow' ? lanes.slow : undefined
+  return {
+    __proto__: null,
+    include: selected?.map(glob =>
+      glob.endsWith('/**') ? `${glob}/*.test.{mts,ts}` : glob,
+    ) ?? ['test/**/*.test.{mts,ts}'],
+    exclude: lane === 'fast' ? [...lanes.mid, ...lanes.slow] : [],
+  }
+}

packages/cli/scripts/test-wrapper.mts

@@ -20,8 +20,10 @@ import { isWin32 } from '@socketsecurity/lib-stable/constants/platform'
 import { getDefaultLogger } from '@socketsecurity/lib-stable/logger/default'
 import { spawn } from '@socketsecurity/lib-stable/process/spawn/child'
 
+import { REPO_ROOT } from '../../../scripts/fleet/paths.mts'
 import { EnvironmentVariables } from './environment-variables.mts'
 import { loadEnvFile } from './util/load-env.mts'
+import { resolvePackageTestScope } from './test-lanes.mts'
 
 const logger = getDefaultLogger()
 const __dirname = path.dirname(fileURLToPath(import.meta.url))
@@ -71,6 +73,9 @@ async function main() {
       args = args.slice(1)
     }
 
+    const scope = resolvePackageTestScope(args, REPO_ROOT)
+    args = scope.args
+
     // Check for and warn about environment variables that can cause snapshot mismatches.
     // These are all aliases for the Socket API token that should not be set during tests.
     const problematicEnvVars = [
@@ -98,6 +103,7 @@ async function main() {
 
     const spawnEnv = {
       ...process.env,
+      FLEET_LANE: scope.lane,
       // Increase Node.js heap size to prevent out of memory errors.
       // Use 8GB in CI, 4GB locally.
       // Add --max-semi-space-size for better GC with RegExp-heavy tests.
@@ -147,6 +153,7 @@ async function main() {
     // On Windows, .cmd files need shell: true.
     const spawnOptions = {
       cwd: rootPath,
+      timeout: scope.timeout,
       env: {
         ...testEnv,
         ...spawnEnv,

packages/cli/test/unit/scripts/test-lanes.test.mts

@@ -0,0 +1,92 @@
+import { expect, it } from 'vitest'
+
+import {
+  coverBudgetMs,
+  laneBudgetMs,
+} from '../../../../../scripts/fleet/constants/test-budget.mts'
+import {
+  packageTestGlobs,
+  resolvePackageTestScope,
+  selectPackageTestGlobs,
+} from '../../../scripts/test-lanes.mts'
+
+it('defaults bare package tests to the fast budget', () => {
+  const scope = resolvePackageTestScope([], '/example/repo')
+  expect(scope).toEqual({
+    args: [],
+    lane: 'fast',
+    timeout: laneBudgetMs('fast'),
+  })
+})
+
+it('keeps explicit files and all-unit runs outside lane filtering', () => {
+  for (const args of [
+    ['--all'],
+    ['test/unit/example.test.mts'],
+    ['--coverage'],
+  ]) {
+    const scope = resolvePackageTestScope(args, '/example/repo')
+    expect(scope.lane).toBeUndefined()
+    expect(scope.timeout).toBe(coverBudgetMs())
+    expect(scope.args).toEqual(args.filter(arg => arg !== '--all'))
+  }
+})
+
+it('consumes lane flags and preserves Vitest arguments', () => {
+  expect(
+    resolvePackageTestScope(
+      ['--lane', 'mid', '--reporter=json'],
+      '/example/repo',
+    ),
+  ).toEqual({
+    args: ['--reporter=json'],
+    lane: 'mid',
+    timeout: laneBudgetMs('mid'),
+  })
+  expect(
+    resolvePackageTestScope(
+      ['--lane=slow', '--config', 'vitest.integration.config.mts'],
+      '/example/repo',
+    ),
+  ).toEqual({
+    args: ['--config', 'vitest.integration.config.mts'],
+    lane: 'slow',
+    timeout: laneBudgetMs('slow'),
+  })
+  expect(() =>
+    resolvePackageTestScope(['--lane=unknown'], '/example/repo'),
+  ).toThrow()
+})
+
+it('normalizes member globs and excludes other packages', () => {
+  expect(
+    packageTestGlobs(
+      [
+        'packages\\cli\\test\\unit\\commands\\**',
+        'packages/cli/test/unit/meow.test.mts',
+        'packages/other/test/unit/example.test.mts',
+        'packages/cli-other/test/unit/example.test.mts',
+      ],
+      'packages\\cli',
+    ),
+  ).toEqual(['test/unit/commands/**', 'test/unit/meow.test.mts'])
+})
+
+it('partitions fast and mid while leaving unscoped coverage complete', () => {
+  const lanes = {
+    mid: ['test/unit/commands/**'],
+    slow: ['test/integration/**'],
+  }
+  expect(selectPackageTestGlobs('fast', lanes)).toEqual({
+    include: ['test/**/*.test.{mts,ts}'],
+    exclude: [...lanes.mid, ...lanes.slow],
+  })
+  expect(selectPackageTestGlobs('mid', lanes)).toEqual({
+    include: ['test/unit/commands/**/*.test.{mts,ts}'],
+    exclude: [],
+  })
+  expect(selectPackageTestGlobs(undefined, lanes)).toEqual({
+    include: ['test/**/*.test.{mts,ts}'],
+    exclude: [],
+  })
+})

packages/cli/vitest.config.mts

@@ -4,9 +4,15 @@ import path from 'node:path'
 import { fileURLToPath } from 'node:url'
 
 import { defineConfig } from 'vitest/config'
+import { REPO_ROOT } from '../../scripts/fleet/paths.mts'
 
 import { vitiatePlugin } from '@vitiate/core/plugin'
 
+import {
+  readPackageTestLanes,
+  selectPackageTestGlobs,
+} from './scripts/test-lanes.mts'
+
 // The vitiate coverage-guided fuzz lane (scripts/repo/fuzz.mts) sets
 // VITIATE_FUZZ=1 and runs `vitest run` against THIS auto-discovered config. In
 // a monorepo package that already owns its vitest.config.mts we can't drop a
@@ -103,6 +109,11 @@ export function getMaxThreads(): number {
   return os.cpus().length
 }
 
+const laneGlobs = selectPackageTestGlobs(
+  process.env['FLEET_LANE'],
+  readPackageTestLanes(REPO_ROOT, 'packages/cli'),
+)
+
 const normalConfig = defineConfig({
   resolve: {
     preserveSymlinks: false,
@@ -125,8 +136,9 @@ const normalConfig = defineConfig({
       NO_COLOR: '',
       TZ: 'UTC',
     },
-    include: ['test/**/*.test.{mts,ts}'],
+    include: laneGlobs.include,
     exclude: [
+      ...laneGlobs.exclude,
       '**/node_modules/**',
       '**/dist/**',
       '**/.{idea,git,cache,output,temp}/**',