From 5a22ffe74fec531cb3f9d7db520a228374dc501c Mon Sep 17 00:00:00 2001
From: Leonid Shevtsov <leonid@shevtsov.me>
Date: Tue, 7 Feb 2017 12:30:08 +0100
Subject: [PATCH] Add option to split files by line count

---
 line_count.go  | 44 ++++++++++++++++++++++++++++++++++++++++++++
 split_tests.go | 31 +++++++++++++++++++------------
 2 files changed, 63 insertions(+), 12 deletions(-)
 create mode 100644 line_count.go

diff --git a/line_count.go b/line_count.go
new file mode 100644
index 0000000..06b5587
--- /dev/null
+++ b/line_count.go
@@ -0,0 +1,44 @@
+package main
+
+import (
+	"bytes"
+	"io"
+	"os"
+)
+
+func estimateFileTimesByLineCount(currentFileSet map[string]bool, fileTimes map[string]float64) {
+	for fileName := range currentFileSet {
+		file, err := os.Open(fileName)
+		if err != nil {
+			printMsg("failed to count lines in file %s: %v\n", file, err)
+			continue
+		}
+		defer file.Close()
+		lineCount, err := lineCounter(file)
+		if err != nil {
+			printMsg("failed to count lines in file %s: %v\n", file, err)
+			continue
+		}
+		fileTimes[fileName] = float64(lineCount)
+	}
+}
+
+// Credit to http://stackoverflow.com/a/24563853/6678
+func lineCounter(r io.Reader) (int, error) {
+	buf := make([]byte, 32*1024)
+	count := 0
+	lineSep := []byte{'\n'}
+
+	for {
+		c, err := r.Read(buf)
+		count += bytes.Count(buf[:c], lineSep)
+
+		switch {
+		case err == io.EOF:
+			return count, nil
+
+		case err != nil:
+			return count, err
+		}
+	}
+}
diff --git a/split_tests.go b/split_tests.go
index 9a2f573..8d1b846 100644
--- a/split_tests.go
+++ b/split_tests.go
@@ -12,6 +12,7 @@ import (
 
 var useCircleCI bool
 var useJUnitXML bool
+var useLineCount bool
 var junitXMLPath string
 var testFilePattern = ""
 var circleCIProjectPrefix = ""
@@ -56,7 +57,9 @@ func addNewFiles(fileTimes map[string]float64, currentFileSet map[string]bool) {
 		if _, isSet := fileTimes[file]; isSet {
 			continue
 		}
-		printMsg("missing file time for %s\n", file)
+		if useCircleCI || useJUnitXML {
+			printMsg("missing file time for %s\n", file)
+		}
 		fileTimes[file] = averageFileTime
 	}
 }
@@ -94,12 +97,14 @@ func parseFlags() {
 	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(&circleCIProjectPrefix, "circleci-project", "", "CircleCI project name (e.g. github/leonid-shevtsov/split_tests) - 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)")
 
+	flag.BoolVar(&useLineCount, "line-count", false, "Use line count to estimate test times")
+
 	var showHelp bool
 	flag.BoolVar(&showHelp, "help", false, "Show this help text")
 
@@ -126,7 +131,6 @@ func parseFlags() {
 	}
 
 	useCircleCI = circleCIAPIKey != ""
-	showHelp = showHelp || !useCircleCI && !useJUnitXML
 
 	if showHelp {
 		printMsg("Splits test files into containers of even duration\n\n")
@@ -144,25 +148,28 @@ func parseFlags() {
 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
 	}
 
+	fileTimes := make(map[string]float64)
+	if useLineCount {
+		estimateFileTimesByLineCount(currentFileSet, fileTimes)
+	} else if useJUnitXML {
+		getFileTimesFromJUnitXML(fileTimes)
+	} else if useCircleCI {
+		getFileTimesFromCircleCI(fileTimes)
+	}
+
 	removeDeletedFiles(fileTimes, currentFileSet)
 	addNewFiles(fileTimes, currentFileSet)
 
 	buckets, bucketTimes := splitFiles(fileTimes, splitTotal)
-	printMsg("expected test time: %0.1fs\n", bucketTimes[splitIndex])
+	if useCircleCI || useJUnitXML {
+		printMsg("expected test time: %0.1fs\n", bucketTimes[splitIndex])
+	}
 
 	fmt.Println(strings.Join(buckets[splitIndex], " "))
 }