diff --git a/scripts/lib/extension-test-plan.mjs b/scripts/lib/extension-test-plan.mjs index 9139228ce2d..191ef850627 100644 --- a/scripts/lib/extension-test-plan.mjs +++ b/scripts/lib/extension-test-plan.mjs @@ -21,6 +21,13 @@ import { listAvailableExtensionIds } from "./changed-extensions.mjs"; const repoRoot = path.resolve(import.meta.dirname, "..", ".."); export const DEFAULT_EXTENSION_TEST_SHARD_COUNT = 6; +const EXTENSION_TEST_COST_MULTIPLIERS = { + "test/vitest/vitest.extension-feishu.config.ts": 1.6, + "test/vitest/vitest.extension-msteams.config.ts": 1.6, + // This shared config is comparatively cheap per file, so raw file count + // overstates its real wall-clock cost during CI shard planning. + "test/vitest/vitest.extensions.config.ts": 0.45, +}; function normalizeRelative(inputPath) { return inputPath.split(path.sep).join("/"); @@ -53,6 +60,11 @@ function countTestFiles(rootPath) { return total; } +function estimatePlanCost(config, testFileCount) { + const multiplier = EXTENSION_TEST_COST_MULTIPLIERS[config] ?? 1; + return Math.max(1, Math.ceil(testFileCount * multiplier)); +} + function resolveExtensionDirectory(targetArg, cwd = process.cwd()) { if (targetArg) { const asGiven = path.resolve(cwd, targetArg); @@ -152,9 +164,11 @@ export function resolveExtensionTestPlan(params = {}) { (sum, root) => sum + countTestFiles(path.join(repoRoot, root)), 0, ); + const estimatedCost = estimatePlanCost(config, testFileCount); return { config, + estimatedCost, extensionDir: relativeExtensionDir, extensionId, hasTests: testFileCount > 0, @@ -171,11 +185,13 @@ function mergeTestPlans(plans) { config: plan.config, extensionIds: [], roots: [], + estimatedCost: 0, testFileCount: 0, }; current.extensionIds.push(plan.extensionId); current.roots.push(...plan.roots); + current.estimatedCost += plan.estimatedCost; current.testFileCount += plan.testFileCount; groupsByConfig.set(plan.config, current); } @@ -193,6 +209,7 @@ function mergeTestPlans(plans) { extensionIds: plans .map((plan) => plan.extensionId) .toSorted((left, right) => left.localeCompare(right)), + estimatedCost: plans.reduce((sum, plan) => sum + plan.estimatedCost, 0), hasTests: plans.length > 0, planGroups, testFileCount: plans.reduce((sum, plan) => sum + plan.testFileCount, 0), @@ -215,6 +232,9 @@ function pickLeastLoadedShard(shards) { return index; } const best = shards[bestIndex]; + if (shard.estimatedCost !== best.estimatedCost) { + return shard.estimatedCost < best.estimatedCost ? index : bestIndex; + } if (shard.testFileCount !== best.testFileCount) { return shard.testFileCount < best.testFileCount ? index : bestIndex; } @@ -233,6 +253,9 @@ export function createExtensionTestShards(params = {}) { .map((extensionId) => resolveExtensionTestPlan({ cwd, targetArg: extensionId })) .filter((plan) => plan.hasTests) .toSorted((left, right) => { + if (left.estimatedCost !== right.estimatedCost) { + return right.estimatedCost - left.estimatedCost; + } if (left.testFileCount !== right.testFileCount) { return right.testFileCount - left.testFileCount; } @@ -241,6 +264,7 @@ export function createExtensionTestShards(params = {}) { const effectiveShardCount = Math.min(shardCount, Math.max(1, plans.length)); const shards = Array.from({ length: effectiveShardCount }, () => ({ + estimatedCost: 0, plans: [], testFileCount: 0, })); @@ -248,6 +272,7 @@ export function createExtensionTestShards(params = {}) { for (const plan of plans) { const targetIndex = pickLeastLoadedShard(shards); shards[targetIndex].plans.push(plan); + shards[targetIndex].estimatedCost += plan.estimatedCost; shards[targetIndex].testFileCount += plan.testFileCount; } diff --git a/test/scripts/test-extension.test.ts b/test/scripts/test-extension.test.ts index a84235dd023..f4c24330182 100644 --- a/test/scripts/test-extension.test.ts +++ b/test/scripts/test-extension.test.ts @@ -282,96 +282,112 @@ describe("scripts/test-extension.mjs", () => { expect(batch.planGroups).toEqual([ { config: "test/vitest/vitest.extension-acpx.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["acpx"], roots: [bundledPluginRoot("acpx")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-bluebubbles.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["bluebubbles"], roots: [bundledPluginRoot("bluebubbles")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-channels.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["line", "slack"], roots: [bundledPluginRoot("slack"), bundledPluginRoot("line")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-diffs.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["diffs"], roots: [bundledPluginRoot("diffs")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-feishu.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["feishu"], roots: [bundledPluginRoot("feishu")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-irc.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["irc"], roots: [bundledPluginRoot("irc")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-matrix.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["matrix"], roots: [bundledPluginRoot("matrix")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-mattermost.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["mattermost"], roots: [bundledPluginRoot("mattermost")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-memory.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["memory-core"], roots: [bundledPluginRoot("memory-core")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-msteams.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["msteams"], roots: [bundledPluginRoot("msteams")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-providers.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["openai"], roots: [bundledPluginRoot("openai")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-telegram.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["telegram"], roots: [bundledPluginRoot("telegram")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-voice-call.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["voice-call"], roots: [bundledPluginRoot("voice-call")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-whatsapp.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["whatsapp"], roots: [bundledPluginRoot("whatsapp")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extension-zalo.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["zalo", "zalouser"], roots: [bundledPluginRoot("zalo"), bundledPluginRoot("zalouser")], testFileCount: expect.any(Number), }, { config: "test/vitest/vitest.extensions.config.ts", + estimatedCost: expect.any(Number), extensionIds: ["firecrawl"], roots: [bundledPluginRoot("firecrawl")], testFileCount: expect.any(Number), @@ -379,7 +395,7 @@ describe("scripts/test-extension.mjs", () => { ]); }); - it("balances extension test shards by test file count", () => { + it("balances extension test shards by estimated CI cost", () => { const shards = createExtensionTestShards({ cwd: process.cwd(), shardCount: DEFAULT_EXTENSION_TEST_SHARD_COUNT, @@ -402,8 +418,15 @@ describe("scripts/test-extension.mjs", () => { ); expect(assigned).toHaveLength(expected.length); - const totals = shards.map((shard) => shard.testFileCount); + const totals = shards.map((shard) => shard.estimatedCost); expect(Math.max(...totals) - Math.min(...totals)).toBeLessThanOrEqual(1); + + const msTeamsShardIndex = shards.findIndex((shard) => shard.extensionIds.includes("msteams")); + const feishuShardIndex = shards.findIndex((shard) => shard.extensionIds.includes("feishu")); + + expect(msTeamsShardIndex).toBeGreaterThanOrEqual(0); + expect(feishuShardIndex).toBeGreaterThanOrEqual(0); + expect(msTeamsShardIndex).not.toBe(feishuShardIndex); }); it("treats extensions without tests as a no-op by default", () => {