Initial commit

Author: leonid-shevtsov

.gitignore

@@ -0,0 +1,2 @@
+build
+vendor

Makefile

@@ -0,0 +1,16 @@
+release: release-linux release-osx
+
+release-linux:
+	$(call release-base,linux)
+
+release-osx:
+	$(call release-base,darwin)
+
+
+release-base = \
+	mkdir -p build/$(1); \
+	GOOS=$(1) GOARCH=amd64 go build -ldflags="-s -w" -o build/$(1)/split_tests; \
+	tar czf build/split_tests-$(1)-$(shell git rev-parse --short HEAD).tar.gz --directory=build/$(1) .
+
+clean:
+	rm -rf build

circleci.go

@@ -0,0 +1,71 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"path"
+)
+
+func getCircleAPIJSON(url string, destination interface{}) {
+	client := &http.Client{}
+	req, err := http.NewRequest("GET", url, nil)
+	req.Header.Add("Accept", "application/json")
+	resp, err := client.Do(req)
+	if err != nil {
+		fatalMsg("error calling CircleCI API at %v: %v", url, err)
+	}
+	defer resp.Body.Close()
+	decoder := json.NewDecoder(resp.Body)
+	err = decoder.Decode(destination)
+	if err != nil {
+		fatalMsg("error parsing CircleCI JSON at %v: %v", url, err)
+	}
+}
+
+type circleCIBranchList []struct {
+	BuildNum int `json:"build_num"`
+}
+
+type circleCITestResults struct {
+	Tests []struct {
+		File    string  `json:"file"`
+		RunTime float64 `json:"run_time"`
+	} `json:"tests"`
+}
+
+func circleCIAPIURL() string {
+	return fmt.Sprintf("https://circleci.com/api/v1.1/project/%s", circleCIProjectPrefix)
+}
+
+func getCircleCIBranchBuilds(branchName string) circleCIBranchList {
+	buildsURL := fmt.Sprintf("%s/tree/%s?filter=successful&circle-token=%s", circleCIAPIURL(), branchName, circleCIAPIKey)
+	var branchList circleCIBranchList
+	getCircleAPIJSON(buildsURL, &branchList)
+	return branchList
+}
+
+func getCircleCITestResults(buildNum int) circleCITestResults {
+	testResultsURL := fmt.Sprintf("%s/%d/tests?circle-token=%s", circleCIAPIURL(), buildNum, circleCIAPIKey)
+	var testResults circleCITestResults
+	getCircleAPIJSON(testResultsURL, &testResults)
+	return testResults
+}
+
+func getFileTimesFromCircleCI(fileTimes map[string]float64) {
+	builds := getCircleCIBranchBuilds(circleCIBranchName)
+	if len(builds) == 0 {
+		builds = getCircleCIBranchBuilds("master")
+	}
+	if len(builds) > 0 {
+		buildNum := builds[0].BuildNum
+		printMsg("using test timings from CircleCI build %d\n", buildNum)
+
+		testResults := getCircleCITestResults(buildNum)
+
+		for _, test := range testResults.Tests {
+			filePath := path.Clean(test.File)
+			fileTimes[filePath] += test.RunTime
+		}
+	}
+}

glide.lock

@@ -0,0 +1,3 @@
+package: github.com/leonid-shevtsov/circleci-test-balancer
+import:
+- package: github.com/mattn/go-zglob

glide.yaml

@@ -0,0 +1,48 @@
+package main
+
+import (
+	"encoding/xml"
+	"io"
+	"os"
+	"path"
+)
+
+type junitXML struct {
+	TestCases []struct {
+		File string  `xml:"file,attr"`
+		Time float64 `xml:"time,attr"`
+	} `xml:"testcase"`
+}
+
+func loadJUnitXML(reader io.Reader) *junitXML {
+	var junitXML junitXML
+
+	decoder := xml.NewDecoder(reader)
+	err := decoder.Decode(&junitXML)
+	if err != nil {
+		fatalMsg("failed to parse junit xml: %v\n", err)
+	}
+
+	return &junitXML
+}
+
+func getFileTimesFromJUnitXML(fileTimes map[string]float64) {
+	var source io.Reader
+	if junitXMLPath != "" {
+		file, err := os.Open(junitXMLPath)
+		if err != nil {
+			fatalMsg("failed to open junit xml: %v\n", err)
+		}
+		defer file.Close()
+		printMsg("using test times from JUnit report %s\n", junitXMLPath)
+		source = file
+	} else {
+		printMsg("using test times from JUnit report at stdin\n")
+		source = os.Stdin
+	}
+	junitXML := loadJUnitXML(source)
+	for _, testCase := range junitXML.TestCases {
+		filePath := path.Clean(testCase.File)
+		fileTimes[filePath] += testCase.Time
+	}
+}

junit.go

@@ -0,0 +1,168 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"os"
+	"strconv"
+	"strings"
+
+	zglob "github.com/mattn/go-zglob"
+)
+
+var useCircleCI bool
+var useJUnitXML bool
+var junitXMLPath string
+var testFilePattern = ""
+var circleCIProjectPrefix = ""
+var circleCIBranchName = os.Getenv("CIRCLE_BRANCH")
+var splitIndex, _ = strconv.Atoi(os.Getenv("CIRCLE_NODE_INDEX"))
+var splitTotal, _ = strconv.Atoi(os.Getenv("CIRCLE_NODE_TOTAL"))
+var circleCIAPIKey = os.Getenv("CIRCLECI_API_KEY")
+
+func printMsg(msg string, args ...interface{}) {
+	if len(args) == 0 {
+		fmt.Fprint(os.Stderr, msg)
+	} else {
+		fmt.Fprintf(os.Stderr, msg, args...)
+	}
+}
+
+func fatalMsg(msg string, args ...interface{}) {
+	printMsg(msg, args...)
+	os.Exit(1)
+}
+
+func removeDeletedFiles(fileTimes map[string]float64, currentFileSet map[string]bool) {
+	for file := range fileTimes {
+		if !currentFileSet[file] {
+			delete(fileTimes, file)
+		}
+	}
+}
+
+func addNewFiles(fileTimes map[string]float64, currentFileSet map[string]bool) {
+	averageFileTime := 0.0
+	if len(fileTimes) > 0 {
+		for _, time := range fileTimes {
+			averageFileTime += time
+		}
+		averageFileTime /= float64(len(fileTimes))
+	} else {
+		averageFileTime = 1.0
+	}
+
+	for file := range currentFileSet {
+		if _, isSet := fileTimes[file]; isSet {
+			continue
+		}
+		printMsg("missing file time for %s\n", file)
+		fileTimes[file] = averageFileTime
+	}
+}
+
+func splitFiles(fileTimes map[string]float64, splitTotal int) ([][]string, []float64) {
+	buckets := make([][]string, splitTotal)
+	bucketTimes := make([]float64, splitTotal)
+	for len(fileTimes) > 0 {
+		// find file with max weight
+		maxFile := ""
+		for file, time := range fileTimes {
+			if len(maxFile) == 0 || time > fileTimes[maxFile] {
+				maxFile = file
+			}
+		}
+		// find bucket with min weight
+		minBucket := 0
+		for bucket := 1; bucket < splitTotal; bucket++ {
+			if bucketTimes[bucket] < bucketTimes[minBucket] {
+				minBucket = bucket
+			}
+		}
+		buckets[minBucket] = append(buckets[minBucket], maxFile)
+		bucketTimes[minBucket] += fileTimes[maxFile]
+		delete(fileTimes, maxFile)
+	}
+
+	return buckets, bucketTimes
+}
+
+func parseFlags() {
+	flag.StringVar(&testFilePattern, "glob", "spec/**/*_spec.rb", "Glob pattern to find test files")
+
+	flag.IntVar(&splitIndex, "split-index", -1, "This test container's index (or set CIRCLE_NODE_INDEX)")
+	flag.IntVar(&splitTotal, "split-total", -1, "Total number of containers (or set CIRCLE_NODE_TOTAL)")
+
+	flag.StringVar(&circleCIAPIKey, "circleci-key", "", "CircleCI API key (or set CIRCLECI_API_KEY environment variable) - required to use CircleCI")
+	flag.StringVar(&circleCIProjectPrefix, "circleci-project", "", "CircleCI project name (e.g. github/leonid-shevtsov/circleci-test-balancer) - required to use CircleCI")
+	flag.StringVar(&circleCIBranchName, "circleci-branch", "", "Current branch for CircleCI (or set CIRCLE_BRANCH) - required to use CircleCI")
+
+	flag.BoolVar(&useJUnitXML, "junit", false, "Use a JUnit XML report for test times")
+	flag.StringVar(&junitXMLPath, "junit-path", "", "Path to a JUnit XML report (leave empty to read from stdin)")
+
+	var showHelp bool
+	flag.BoolVar(&showHelp, "help", false, "Show this help text")
+
+	flag.Parse()
+
+	var err error
+	if circleCIAPIKey == "" {
+		circleCIAPIKey = os.Getenv("CIRCLECI_API_KEY")
+	}
+	if circleCIBranchName == "" {
+		circleCIBranchName = os.Getenv("CIRCLE_BRANCH")
+	}
+	if splitTotal == -1 {
+		splitTotal, err = strconv.Atoi(os.Getenv("CIRCLE_NODE_TOTAL"))
+		if err != nil {
+			splitIndex = -1
+		}
+	}
+	if splitIndex == -1 {
+		splitIndex, err = strconv.Atoi(os.Getenv("CIRCLE_NODE_INDEX"))
+		if err != nil {
+			splitIndex = -1
+		}
+	}
+
+	useCircleCI = circleCIAPIKey != ""
+	showHelp = showHelp || !useCircleCI && !useJUnitXML
+
+	if showHelp {
+		printMsg("Splits test files into containers of even duration\n\n")
+		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()
+
+	fileTimes := make(map[string]float64)
+	if useJUnitXML {
+		getFileTimesFromJUnitXML(fileTimes)
+	} else if useCircleCI {
+		getFileTimesFromCircleCI(fileTimes)
+	}
+
+	currentFiles, _ := zglob.Glob(testFilePattern)
+	currentFileSet := make(map[string]bool)
+
+	for _, file := range currentFiles {
+		currentFileSet[file] = true
+	}
+
+	removeDeletedFiles(fileTimes, currentFileSet)
+	addNewFiles(fileTimes, currentFileSet)
+
+	buckets, bucketTimes := splitFiles(fileTimes, splitTotal)
+	printMsg("expected test time: %0.1fs\n", bucketTimes[splitIndex])
+
+	fmt.Println(strings.Join(buckets[splitIndex], " "))
+}