Support test partitioning (#28)

* Support test partitioning

* Test fix

* Update README.md

Co-authored-by: appleby <[email protected]>

---------

Co-authored-by: appleby <[email protected]>

Author: camsaul

README.md

@@ -202,11 +202,33 @@ they live in is loaded; this may be affected by `:only` options passed to the te
 
 Return values of methods are ignored; they are done purely for side effects.
 
+## Partitioning tests
+
+You can divide a test suite into multiple partitions using the `:partition/total` and `:partition/index` keys. This is
+an easy way to speed up CI by diving large test suites into multiple jobs.
+
+```
+clj -X:test '{:partition/total 10, :partition/index 8}'
+...
+Running tests in partition 9 of 10 (575 tests of 5753)...
+Finding tests took 46.6 s.
+Running 575 tests
+...
+```
+
+`:partition/index` is zero-based, e.g. if you have ten partitions (`:partiton/total 10`) then the first partition is `0` and
+the last is `9`.
+
+Tests are partitioned at the `deftest` level after all tests are found the usual way -- all namespaces that would be
+loaded if you were running the entire test suite are still loaded. Partitions are split as evenly as possible, but
+tests are guaranteed to be split deterministically into exactly the number of partitions you asked for.
+
+
 ## Additional options
 
 All other options are passed directly to [Eftest](https://github.com/weavejester/eftest); refer to its documentation
 for more information.
 
 ```
-clj -X:test :fail-fast? true
+clj -X:test '{:fail-fast? true}'
 ```

src/mb/hawk/core.clj

@@ -2,6 +2,7 @@
   (:require
    [clojure.java.classpath :as classpath]
    [clojure.java.io :as io]
+   [clojure.math :as math]
    [clojure.pprint :as pprint]
    [clojure.set :as set]
    [clojure.string :as str]
@@ -120,14 +121,56 @@
   [_nil options]
   (find-tests (classpath/system-classpath) options))
 
+(defn- partition-all-into-n-partitions
+  "Split sequence `xs` into `num-partitions` as equally as possible. Guaranteed to return `num-partitions`. This custom
+  function is used instead of [[partition-all]] or whatever because we want to make sure every partition gets tests,
+  even with weird combinations like 4 tests with 3 partitions or 29 tests with 10 partitions."
+  [num-partitions xs]
+  {:post [(= (count %) num-partitions)]}
+  ;; make sure the partitioning is deterministic -- `xs` should always come back in the same order but we should sort
+  ;; just to be safe.
+  (let [xs             (sort-by str xs)
+        partition-size (/ (count xs) num-partitions)]
+    (into []
+          (comp (map-indexed (fn [i x]
+                               [(long (math/floor (/ i partition-size))) x]))
+                (partition-by first)
+                (map (fn [partition]
+                       (map second partition))))
+          xs)))
+
+(defn- partition-tests [tests {num-partitions :partition/total, partition-index :partition/index, :as _options}]
+  (if (or num-partitions partition-index)
+    (do
+      (assert (and num-partitions partition-index)
+              ":partition/total and :partition/index must be set together")
+      (assert (pos-int? num-partitions)
+              "Invalid :partition/total - must be a positive integer")
+      (assert (<= num-partitions (count tests))
+              "Invalid :partition/total - cannot have more partitions than number of tests")
+      (assert (int? partition-index)
+              "Invalid :partition/index - must be an integer")
+      (assert (<= 0 partition-index (dec num-partitions))
+              (format "Invalid :partition/index - must be between 0 and %d" (dec num-partitions)))
+      (let [partitions (partition-all-into-n-partitions num-partitions tests)
+            partition  (nth partitions partition-index)]
+        (printf "Running tests in partition %d of %d (%d tests of %d)...\n"
+                (inc partition-index)
+                num-partitions
+                (count partition)
+                (count tests))
+        partition))
+    tests))
+
 (defn find-tests-with-options
   "Find tests using the options map as passed to `clojure -X`."
   [{:keys [only], :as options}]
   (println "Running tests with options" (pr-str options))
   (when only
     (println "Running tests in" (pr-str only)))
   (let [start-time-ms (System/currentTimeMillis)
-        tests         (find-tests only options)]
+        tests         (-> (find-tests only options)
+                          (partition-tests options))]
     (printf "Finding tests took %s.\n" (u/format-milliseconds (- (System/currentTimeMillis) start-time-ms)))
     (println "Running" (count tests) "tests")
     tests))

test/mb/hawk/core_test.clj

@@ -58,3 +58,48 @@
     {:exclude-tags [:exclude-this-test]}
     {:exclude-tags #{:exclude-this-test}}
     {:exclude-tags [:exclude-this-test :another/tag]}))
+
+(deftest ^:parallel partition-tests-test
+  (are [i expected] (= expected
+                       (#'hawk/partition-tests
+                        (range 4)
+                        {:partition/index i, :partition/total 3}))
+    0 [0 1]
+    1 [2]
+    2 [3])
+  (are [i expected] (= expected
+                       (#'hawk/partition-tests
+                        (range 5)
+                        {:partition/index i, :partition/total 3}))
+    0 [0 1]
+    1 [2 3]
+    2 [4]))
+
+(deftest ^:parallel partition-tests-determinism-test
+  (testing "partitioning should be deterministic even if tests come back in a non-deterministic order for some reason"
+    (are [i expected] (= expected
+                         (#'hawk/partition-tests
+                          (shuffle (map #(format "%02d" %) (range 26)))
+                          {:partition/index i, :partition/total 10}))
+      0 ["00" "01" "02"]
+      1 ["03" "04" "05"]
+      2 ["06" "07"]
+      3 ["08" "09" "10"]
+      4 ["11" "12"]
+      5 ["13" "14" "15"]
+      6 ["16" "17" "18"]
+      7 ["19" "20"]
+      8 ["21" "22" "23"]
+      9 ["24" "25"])))
+
+(deftest ^:parallel partition-test
+  (are [index expected] (= expected
+                           (hawk/find-tests-with-options {:only            `[find-tests-test
+                                                                             exclude-tags-test
+                                                                             partition-tests-test
+                                                                             partition-test]
+                                                          :partition/index index
+                                                          :partition/total 3}))
+    0 [#'exclude-tags-test #'find-tests-test]
+    1 [#'partition-test]
+    2 [#'partition-tests-test]))