Coverage for scripts / split_tests.py: 100.00%

1#!/usr/bin/env python3 

2"""Partition the core pytest suite into balanced shards for parallel CI jobs. 

3 

4``unit:pytest:core`` ran the whole suite in one job and had grown to a steady 

514-15 minutes against a 20-minute timeout — close enough that a normal increase 

6in test count tipped it over. Splitting the run across jobs restores headroom 

7and halves the time contributors wait for a result. 

8 

9The split is computed at run time rather than checked in, so it cannot drift as 

10test files are added, renamed, or deleted. Collection is asked for the test count 

11of every file, files are sorted heaviest first, and each is assigned to whichever 

12shard currently holds the fewest tests. That is a greedy bin-pack: not optimal, 

13but it keeps shards within a few percent of each other and is deterministic, so 

14a rerun of the same commit produces the same partition. 

15 

16Test *count* is a proxy for test *duration*. It is a good enough proxy here 

17because the slow, uneven work — CDK synthesis, cdk-nag, cross-module integration 

18— already lives in its own dedicated jobs and is excluded below. 

19 

20Usage:: 

21 

22 # Files for shard 1 of 2, one per line 

23 python scripts/split_tests.py --shard 1 --of 2 

24 

25 # What the partition looks like, without running anything 

26 python scripts/split_tests.py --of 2 --summary 

27 

28Exits non-zero if collection fails, so a broken suite fails the job rather than 

29silently producing an empty shard. 

30""" 

31 

32from __future__ import annotations 

33 

34import argparse 

35import json 

36import subprocess 

37import sys 

38from pathlib import Path 

39 

40REPO_ROOT = Path(__file__).resolve().parent.parent 

41SUITE_ROOT = "tests" 

42 

43#: Test modules excluded from the core suite because a dedicated CI job runs 

44#: them. Keeping the mapping here — rather than as bare ``--ignore`` flags in 

45#: the workflow — means the two shard jobs cannot disagree about the exclusions, 

46#: and the reason for each is visible next to the path. 

47DEDICATED_JOB_MODULES: dict[str, str] = { 

48 "tests/test_integration.py": "integration:pytest:cross-module", 

49 "tests/test_mcp_integration.py": "integration:mcp:server", 

50 "tests/test_nag_compliance.py": "unit:cdk:nag-compliance", 

51 "tests/test_cdk_synthesis_matrix.py": "unit:cdk:config-matrix", 

52 "tests/test_project_name_scoping.py": "unit:cdk:project-name-scoping", 

53 "tests/test_accelerator_catalog.py": "unit:pytest:core, offline policy step", 

54 "tests/test_accelerator_pools.py": "unit:pytest:core, offline policy step", 

55} 

56 

57 

58def ignore_args() -> list[str]: 

59 """Return the ``--ignore`` flags that carve the core suite out of ``tests/``.""" 

60 return [f"--ignore={path}" for path in sorted(DEDICATED_JOB_MODULES)] 

61 

62 

63def collect_counts() -> dict[str, int]: 

64 """Return ``{test file: number of collected tests}`` for the core suite. 

65 

66 Uses ``pytest --collect-only -q``, which lists one node id per line, so the 

67 counts reflect parametrization rather than the number of ``def test_`` 

68 statements. 

69 """ 

70 result = subprocess.run( 

71 [ 

72 sys.executable, 

73 "-m", 

74 "pytest", 

75 SUITE_ROOT, 

76 # ``-o addopts=`` clears the project's ``addopts`` (which sets -v). 

77 # Without it verbosity nets to 0 and --collect-only prints an 

78 # indented tree instead of one node id per line, which this parser 

79 # cannot read. Overriding rather than adding more -q flags keeps the 

80 # output shape independent of project-level verbosity settings. 

81 "-o", 

82 "addopts=", 

83 "-q", 

84 "--collect-only", 

85 "--no-header", 

86 "-p", 

87 "no:cacheprovider", 

88 *ignore_args(), 

89 ], 

90 cwd=REPO_ROOT, 

91 capture_output=True, 

92 text=True, 

93 check=False, 

94 ) 

95 if result.returncode != 0: 

96 sys.stderr.write(result.stdout) 

97 sys.stderr.write(result.stderr) 

98 raise SystemExit( 

99 f"pytest collection failed with exit code {result.returncode}; " 

100 "cannot partition the suite" 

101 ) 

102 

103 counts: dict[str, int] = {} 

104 for line in result.stdout.splitlines(): 

105 node = line.strip() 

106 if "::" not in node: 

107 continue 

108 path = node.split("::", 1)[0] 

109 if not path.endswith(".py"): 

110 continue 

111 counts[path] = counts.get(path, 0) + 1 

112 

113 if not counts: 

114 raise SystemExit("pytest collection produced no test ids; refusing to shard") 

115 return counts 

116 

117 

118def balance(counts: dict[str, int], shards: int) -> list[list[str]]: 

119 """Greedily bin-pack files into ``shards`` groups of similar test count. 

120 

121 Files are placed heaviest first into the currently lightest shard. Ties are 

122 broken by path so the partition is stable across runs. 

123 """ 

124 if shards < 1: 

125 raise SystemExit("--of must be at least 1") 

126 

127 groups: list[list[str]] = [[] for _ in range(shards)] 

128 totals = [0] * shards 

129 

130 for path, count in sorted(counts.items(), key=lambda item: (-item[1], item[0])): 

131 target = min(range(shards), key=lambda index: (totals[index], index)) 

132 groups[target].append(path) 

133 totals[target] += count 

134 

135 return [sorted(group) for group in groups] 

136 

137 

138def main(argv: list[str] | None = None) -> int: 

139 parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) 

140 parser.add_argument("--of", type=int, required=True, help="Total number of shards.") 

141 parser.add_argument("--shard", type=int, help="Which shard to print (1-based).") 

142 parser.add_argument( 

143 "--summary", 

144 action="store_true", 

145 help="Print the per-shard test counts instead of a file list.", 

146 ) 

147 parser.add_argument("--json", action="store_true", help="Emit the partition as JSON.") 

148 args = parser.parse_args(argv) 

149 

150 counts = collect_counts() 

151 groups = balance(counts, args.of) 

152 

153 if args.json: 

154 print( 

155 json.dumps( 

156 { 

157 "total_tests": sum(counts.values()), 

158 "total_files": len(counts), 

159 "shards": [ 

160 {"shard": i + 1, "tests": sum(counts[p] for p in g), "files": g} 

161 for i, g in enumerate(groups) 

162 ], 

163 }, 

164 indent=2, 

165 ) 

166 ) 

167 return 0 

168 

169 if args.summary: 

170 total = sum(counts.values()) 

171 print(f"{total} tests across {len(counts)} files -> {args.of} shard(s)") 

172 for index, group in enumerate(groups, 1): 

173 shard_total = sum(counts[path] for path in group) 

174 share = (shard_total / total * 100) if total else 0 

175 print(f" shard {index}: {shard_total:5d} tests ({share:5.1f}%) in {len(group)} files") 

176 return 0 

177 

178 if args.shard is None: 

179 raise SystemExit("--shard is required unless --summary or --json is given") 

180 if not 1 <= args.shard <= args.of: 

181 raise SystemExit(f"--shard must be between 1 and {args.of}") 

182 

183 for path in groups[args.shard - 1]: 

184 print(path) 

185 return 0 

186 

187 

188if __name__ == "__main__": 

189 raise SystemExit(main())