stabilize test timings

Author: leonid-shevtsov

README.md

@@ -40,6 +40,18 @@ Or, if it's easier to pipe the report file:
 rspec $(curl http://my.junit.url | split_tests -junit -split-index=$CI_NODE_INDEX -split-total=$CI_NODE_TOTAL)
 ```
 
+#### Stabilizing test timings
+
+Sometimes test timings fluctuate, so always relying on the latest timings might not make the best split. For such cases, `split_tests` has an update mode that will apply a sliding window average between prior timings and current ones:
+
+```
+split_tests -junit-update=old_glob -junit-new=new_glob -junit-out=out.xml
+```
+
+Then you take out.xml and use it for the next test run.
+
+Note that updating also cleans up the files and only preserves one "test case" per file, because that is enough for this tool's purpose.
+
 ### Naive split by line count
 
 If you don't have test times, it might be reasonable for your project to assume runtime proportional to test length.
@@ -95,6 +107,12 @@ $./split_tests -help
         Use a JUnit XML report for test times
   -junit-path string
         Path to a JUnit XML report (leave empty to read from stdin; use glob pattern to load multiple files)
+  -junit-new string
+        Glob pattern for new JUnit XML files (for updating timings with sliding window)
+  -junit-out string
+        Output path for updated JUnit XML file (for updating timings with sliding window)
+  -junit-update string
+        Glob pattern for old JUnit XML files (for updating timings with sliding window)
   -line-count
         Use line count to estimate test times
   -split-index int

junit.go

@@ -36,20 +36,42 @@ func addFileTimesFromIOReader(fileTimes map[string]float64, reader io.Reader) {
 	}
 }
 
-func getFileTimesFromJUnitXML(fileTimes map[string]float64) {
-	if junitXMLPath != "" {
-		filenames, err := doublestar.Glob(junitXMLPath)
+// loadJUnitTimingsFromGlob loads test timings from JUnit XML files matching a glob pattern
+func loadJUnitTimingsFromGlob(globPattern string) map[string]float64 {
+	fileTimes := make(map[string]float64)
+
+	if globPattern == "" {
+		return fileTimes
+	}
+
+	filenames, err := doublestar.Glob(globPattern)
+	if err != nil {
+		fatalMsg("failed to match jUnit filename pattern: %v", err)
+	}
+
+	if len(filenames) == 0 {
+		printMsg("warning: no files matched pattern %s\n", globPattern)
+		return fileTimes
+	}
+
+	for _, junitFilename := range filenames {
+		file, err := os.Open(junitFilename)
 		if err != nil {
-			fatalMsg("failed to match jUnit filename pattern: %v", err)
+			fatalMsg("failed to open junit xml: %v\n", err)
 		}
-		for _, junitFilename := range filenames {
-			file, err := os.Open(junitFilename)
-			if err != nil {
-				fatalMsg("failed to open junit xml: %v\n", err)
-			}
-			printMsg("using test times from JUnit report %s\n", junitFilename)
-			addFileTimesFromIOReader(fileTimes, file)
-			file.Close()
+		printMsg("loaded test times from %s\n", junitFilename)
+		addFileTimesFromIOReader(fileTimes, file)
+		file.Close()
+	}
+
+	return fileTimes
+}
+
+func getFileTimesFromJUnitXML(fileTimes map[string]float64) {
+	if junitXMLPath != "" {
+		loadedTimes := loadJUnitTimingsFromGlob(junitXMLPath)
+		for file, time := range loadedTimes {
+			fileTimes[file] += time
 		}
 	} else {
 		printMsg("using test times from JUnit report at stdin\n")

junit_update.go

@@ -0,0 +1,114 @@
+package main
+
+import (
+	"encoding/xml"
+	"os"
+	"strconv"
+)
+
+var junitUpdateOldGlob string
+var junitUpdateNewGlob string
+var junitUpdateOutPath string
+
+const slidingWindowOldWeight = 0.9
+
+// applySlidingWindow applies an exponential moving average to smooth out timing fluctuations
+// Uses exponential moving average: oldWeight * old + (1 - oldWeight) * new
+// This gives more weight to historical data while still incorporating recent changes
+func applySlidingWindow(oldTime, newTime float64) float64 {
+	return slidingWindowOldWeight*oldTime + (1-slidingWindowOldWeight)*newTime
+}
+
+
+// updateJUnitTimings merges old and new JUnit timings using a sliding window algorithm
+func updateJUnitTimings() {
+	if junitUpdateOldGlob == "" || junitUpdateNewGlob == "" || junitUpdateOutPath == "" {
+		fatalMsg("junit-update requires -junit-update, -junit-new, and -junit-out flags\n")
+	}
+
+	// Load old timings
+	oldTimings := loadJUnitTimingsFromGlob(junitUpdateOldGlob)
+	printMsg("loaded %d test files from old timings\n", len(oldTimings))
+
+	// Load new timings
+	newTimings := loadJUnitTimingsFromGlob(junitUpdateNewGlob)
+	printMsg("loaded %d test files from new timings\n", len(newTimings))
+
+	// Merge timings using sliding window algorithm
+	mergedTimings := make(map[string]float64)
+
+	// Process all tests from new timings
+	for file, newTime := range newTimings {
+		oldTime, exists := oldTimings[file]
+		if exists {
+			// Test exists in both: apply sliding window
+			mergedTimings[file] = applySlidingWindow(oldTime, newTime)
+		} else {
+			// Test not in old: use new timing
+			mergedTimings[file] = newTime
+		}
+	}
+
+	// Tests not in new are automatically excluded (not added to mergedTimings)
+
+	printMsg("merged %d test files (removed tests not in new, used sliding window for existing tests)\n", len(mergedTimings))
+
+	// Write output JUnit XML
+	writeJUnitXML(mergedTimings, junitUpdateOutPath)
+	printMsg("wrote updated timings to %s\n", junitUpdateOutPath)
+}
+
+// writeJUnitXML writes test timings to a JUnit XML file
+func writeJUnitXML(timings map[string]float64, outputPath string) {
+	// JUnit XML structure for writing (with testsuite root element)
+	type testCase struct {
+		File string `xml:"file,attr"`
+		Time string `xml:"time,attr"`
+	}
+
+	type testSuite struct {
+		XMLName   xml.Name   `xml:"testsuite"`
+		Name      string     `xml:"name,attr"`
+		Tests     int        `xml:"tests,attr"`
+		TestCases []testCase `xml:"testcase"`
+	}
+
+	// Convert map to slice for consistent output
+	testCases := make([]testCase, 0, len(timings))
+	for file, time := range timings {
+		// Format as decimal without scientific notation
+		timeStr := strconv.FormatFloat(time, 'f', -1, 64)
+		testCases = append(testCases, testCase{
+			File: file,
+			Time: timeStr,
+		})
+	}
+
+	suite := testSuite{
+		Name:      "rspec",
+		Tests:     len(testCases),
+		TestCases: testCases,
+	}
+
+	// Create output file
+	file, err := os.Create(outputPath)
+	if err != nil {
+		fatalMsg("failed to create output file: %v\n", err)
+	}
+	defer file.Close()
+
+	// Write XML header
+	file.WriteString(xml.Header)
+
+	// Create encoder
+	encoder := xml.NewEncoder(file)
+	encoder.Indent("", "  ")
+
+	// Encode XML
+	err = encoder.Encode(suite)
+	if err != nil {
+		fatalMsg("failed to encode JUnit XML: %v\n", err)
+	}
+
+	file.WriteString("\n")
+}

main.go

@@ -82,6 +82,10 @@ func parseFlags() {
 
 	flag.BoolVar(&useLineCount, "line-count", false, "Use line count to estimate test times")
 
+	flag.StringVar(&junitUpdateOldGlob, "junit-update", "", "Glob pattern for old JUnit XML files (for updating timings with sliding window)")
+	flag.StringVar(&junitUpdateNewGlob, "junit-new", "", "Glob pattern for new JUnit XML files (for updating timings with sliding window)")
+	flag.StringVar(&junitUpdateOutPath, "junit-out", "", "Output path for updated JUnit XML file (for updating timings with sliding window)")
+
 	var showHelp bool
 	flag.BoolVar(&showHelp, "help", false, "Show this help text")
 
@@ -116,17 +120,26 @@ func parseFlags() {
 		flag.PrintDefaults()
 		os.Exit(1)
 	}
+
 	if useCircleCI && (circleCIProjectPrefix == "" || circleCIBranchName == "") {
 		fatalMsg("Incomplete CircleCI configuration (set -circleci-key, -circleci-project, and -circleci-branch\n")
 	}
-	if splitTotal == 0 || splitIndex < 0 || splitIndex > splitTotal {
-		fatalMsg("-split-index and -split-total (and environment variables) are missing or invalid\n")
-	}
 }
 
 func main() {
 	parseFlags()
 
+	// If JUnit update mode is enabled, handle it separately and exit
+	if junitUpdateOldGlob != "" || junitUpdateNewGlob != "" || junitUpdateOutPath != "" {
+		updateJUnitTimings()
+		return
+	}
+
+	// Validate split parameters (not needed in update mode)
+	if splitTotal == 0 || splitIndex < 0 || splitIndex > splitTotal {
+		fatalMsg("-split-index and -split-total (and environment variables) are missing or invalid\n")
+	}
+
 	// We are not using filepath.Glob,
 	// because it doesn't support '**' (to match all files in all nested directories)
 	currentFiles, err := doublestar.Glob(testFilePattern)