Skip to content

Commit

Permalink
Code cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@camsaul
camsaul committed Sep 5, 2024
1 parent 8e474cd commit 828d000
Show file tree
Hide file tree
Showing 4 changed files with 237 additions and 174 deletions.
108 changes: 2 additions & 106 deletions src/mb/hawk/core.clj
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
[mb.hawk.init :as hawk.init]
[mb.hawk.junit :as hawk.junit]
[mb.hawk.parallel :as hawk.parallel]
[mb.hawk.partition :as hawk.partition]
[mb.hawk.speak :as hawk.speak]
[mb.hawk.util :as u]))

Expand Down Expand Up @@ -121,111 +122,6 @@
[_nil options]
(find-tests (classpath/system-classpath) options))

(defn- namespace*
"Like [[clojure.core/namespace]] but handles vars."
[x]
(cond
(instance? clojure.lang.Named x) (namespace x)
(var? x) (namespace (symbol x))
:else nil))

(defn- ensure-test-namespaces-are-contiguous
"Make sure `test-vars` have all the tests for each namespace next to one another so when we split we can do so without
splitting in the middle of a namespace. Does not otherwise change the order of the tests or namespaces."
[test-vars]
(let [namespace->sort-position (into {}
(map-indexed
(fn [i nmspace]
[nmspace i]))
(distinct (map namespace* test-vars)))
test-var->sort-position (into {}
(map-indexed
(fn [i varr]
[varr i]))
test-vars)]
(sort-by (juxt #(namespace->sort-position (namespace* %))
test-var->sort-position)
test-vars)))

(defn- make-test-var->partition [num-partitions test-vars]
(let [;; first figure out approximately how big each partition should be.
target-partition-size (/ (count test-vars) num-partitions)
;; then for each test var figure out what partition it would live in ideally if we weren't worried about making
;; sure namespaces are grouped together.
test-var->ideal-partition (into {}
(map-indexed (fn [i test-var]
(let [ideal-partition (long (math/floor (/ i target-partition-size)))]
(assert (<= 0 ideal-partition (dec num-partitions)))
[test-var ideal-partition]))
test-vars))
;; For each namespace figure out how many tests are in each and what the possible partitions we can put that
;; namespace into. For most namespaces there should only be one possible partition but for some the ideal split
;; happens in the middle of the namespace which means we have two possible candidate partitions to put it into.
namespace->num-tests (reduce
(fn [m test-var]
(update m (namespace* test-var) (fnil inc 0)))
{}
test-vars)
namespace->possible-partitions (reduce
(fn [m test-var]
(update m (namespace* test-var) #(conj (set %) (test-var->ideal-partition test-var))))
{}
test-vars)
;; Decide the canonical partition for each namespace. Keep track of how many tests are in each partititon. If
;; there are multiple possible candidate partitions for a namespace, choose the one that has the least tests in
;; it.
namespace->partition (:namespace->partition
(reduce
(fn [m nmspace]
(let [partition (first (sort-by (fn [partition]
(get-in m [:partition->size partition]))
(namespace->possible-partitions nmspace)))]
(-> m
(update-in [:partition->size partition] (fnil + 0) (namespace->num-tests nmspace))
(assoc-in [:namespace->partition nmspace] partition))))
{}
;; process namespaces in the order they appear in test-vars
(distinct (map namespace* test-vars))))]
(fn test-var->partition [test-var]
(get namespace->partition (namespace* test-var)))))

(defn- partition-tests-into-n-partitions
"Split a sequence of `test-vars` into `num-partitions` (returning a map of partition number => sequence of tests).
Attempts to divide tests up into partitions that are as equal as possible, but keeps tests in the same namespace
grouped together."
[num-partitions test-vars]
{:post [(= (count %) num-partitions)]}
(let [test-vars (ensure-test-namespaces-are-contiguous test-vars)
test-var->partition (make-test-var->partition num-partitions test-vars)]
(reduce
(fn [m test-var]
(update m (test-var->partition test-var) #(conj (vec %) test-var)))
(sorted-map)
test-vars)))

(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 [partition-index->tests (partition-tests-into-n-partitions num-partitions tests)
partition (get partition-index->tests 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}]
Expand All @@ -234,7 +130,7 @@
(println "Running tests in" (pr-str only)))
(let [start-time-ms (System/currentTimeMillis)
tests (-> (find-tests only options)
(partition-tests options))]
(hawk.partition/partition-tests options))]
(printf "Finding tests took %s.\n" (u/format-milliseconds (- (System/currentTimeMillis) start-time-ms)))
(println "Running" (count tests) "tests")
tests))
Expand Down
156 changes: 156 additions & 0 deletions src/mb/hawk/partition.clj
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,156 @@
(ns mb.hawk.partition
(:require
[clojure.math :as math]))

(defn- namespace*
"Like [[clojure.core/namespace]] but handles vars."
[x]
(cond
(instance? clojure.lang.Named x) (namespace x)
(var? x) (namespace (symbol x))
:else nil))

(defn- ensure-test-namespaces-are-contiguous
"Make sure `test-vars` have all the tests for each namespace next to one another so when we split we can do so without
splitting in the middle of a namespace. Does not otherwise change the order of the tests or namespaces."
[test-vars]
(let [namespace->sort-position (into {}
(map-indexed
(fn [i nmspace]
[nmspace i]))
(distinct (map namespace* test-vars)))
test-var->sort-position (into {}
(map-indexed
(fn [i varr]
[varr i]))
test-vars)]
(sort-by (juxt #(namespace->sort-position (namespace* %))
test-var->sort-position)
test-vars)))

(defn- namespace->num-tests
"Return a map of
namespace string => number of tests in that namespace"
[test-vars]
(reduce
(fn [m test-var]
(update m (namespace* test-var) (fnil inc 0)))
{}
test-vars))

(defn- test-var->ideal-partition
"Return a map of
test-var => ideal partition number
'Ideal partition number' is the partition it would live in ideally if we weren't worried about making sure namespaces
are grouped together."
[num-partitions test-vars]
(let [target-partition-size (/ (count test-vars) num-partitions)]
(into {}
(map-indexed (fn [i test-var]
(let [ideal-partition (long (math/floor (/ i target-partition-size)))]
(assert (<= 0 ideal-partition (dec num-partitions)))
[test-var ideal-partition]))
test-vars))))

(defn- namespace->possible-partitions
"Return a map of
namespace string => set of possible partition numbers for its tests
For most namespaces there should only be one possible partition but for some the ideal split happens in the middle of
the namespace which means we have two possible candidate partitions to put it into."
[num-partitions test-vars]
(let [test-var->ideal-partition (test-var->ideal-partition num-partitions test-vars)]
(reduce
(fn [m test-var]
(update m (namespace* test-var) #(conj (set %) (test-var->ideal-partition test-var))))
{}
test-vars)))

(defn- namespace->partition
"Return a map of
namespace string => canonical partition number for its tests
If there are multiple possible candidate partitions for a namespace, choose the one that has the least tests in it."
[num-partitions test-vars]
(let [namespace->num-tests (namespace->num-tests test-vars)
namespace->possible-partitions (namespace->possible-partitions num-partitions test-vars)
;; process all the namespaces that have no question about what partition they should go into first so we have as
;; accurate a picture of the size of each partition as possible before dealing with the ambiguous ones
namespaces (distinct (map namespace* test-vars))
multiple-possible-partitions? (fn [nmspace]
(> (count (namespace->possible-partitions nmspace))
1))
namespaces (concat (remove multiple-possible-partitions? namespaces)
(filter multiple-possible-partitions? namespaces))]
;; Keep track of how many tests are in each partition so far
(:namespace->partition
(reduce
(fn [m nmspace]
(let [partition (first (sort-by (fn [partition]
(get-in m [:partition->size partition]))
(namespace->possible-partitions nmspace)))]
(-> m
(update-in [:partition->size partition] (fnil + 0) (namespace->num-tests nmspace))
(assoc-in [:namespace->partition nmspace] partition))))
{}
namespaces))))

(defn- make-test-var->partition
"Return a function with the signature
(f test-var) => partititon-number"
[num-partitions test-vars]
(let [namespace->partition (namespace->partition num-partitions test-vars)]
(fn test-var->partition [test-var]
(get namespace->partition (namespace* test-var)))))

(defn- partition-tests-into-n-partitions
"Split a sequence of `test-vars` into `num-partitions`, returning a map of
partition number => sequence of tests
Attempts to divide tests up into partitions that are as equal as possible, but keeps tests in the same namespace
grouped together."
[num-partitions test-vars]
{:post [(= (count %) num-partitions)]}
(let [test-vars (ensure-test-namespaces-are-contiguous test-vars)
test-var->partition (make-test-var->partition num-partitions test-vars)]
(reduce
(fn [m test-var]
(update m (test-var->partition test-var) #(conj (vec %) test-var)))
(sorted-map)
test-vars)))

(defn- validate-partition-options [tests {num-partitions :partition/total, partition-index :partition/index, :as _options}]
(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))))

(defn partition-tests
"Return only `tests` to run for the current partition (if `:partition/total` and `:partition/index` are specified). If
they are not specified this returns all `tests`."
[tests {num-partitions :partition/total, partition-index :partition/index, :as options}]
(if (or num-partitions partition-index)
(do
(validate-partition-options tests options)
(let [partition-index->tests (partition-tests-into-n-partitions num-partitions tests)
partition (get partition-index->tests 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))
69 changes: 1 addition & 68 deletions test/mb/hawk/core_test.clj
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
(ns ^:exclude-tags-test ^:mic/test mb.hawk.core-test
(:require
[clojure.test :refer :all]
[mb.hawk.core :as hawk]
[mb.hawk.parallel-test]
[mb.hawk.speak-test]))
[mb.hawk.core :as hawk]))

(deftest ^:exclude-this-test find-tests-test
(testing "symbol naming"
Expand Down Expand Up @@ -60,68 +58,3 @@
{:exclude-tags [:exclude-this-test]}
{:exclude-tags #{:exclude-this-test}}
{:exclude-tags [:exclude-this-test :another/tag]}))

(defn- partition-tests* [num-partitions tests]
(into (sorted-map)
(map (fn [i]
[i (#'hawk/partition-tests
tests
{:partition/index i, :partition/total num-partitions})]))
(range num-partitions)))

(deftest ^:parallel partition-tests-test
(is (= '{0 [a/test b/test]
1 [c/test]
2 [d/test]}
(partition-tests* 3 '[a/test b/test c/test d/test])))
(is (= '{0 [a/test b/test]
1 [c/test d/test]
2 [e/test]}
(partition-tests* 3 '[a/test b/test c/test d/test e/test]))))

(deftest ^:parallel partition-tests-evenly-test
(testing "make sure we divide things roughly evenly"
(is (= '{0 [n00/test n01/test n02/test]
1 [n03/test n04/test n05/test]
2 [n06/test n07/test]
3 [n08/test n09/test n10/test]
4 [n11/test n12/test]
5 [n13/test n14/test n15/test]
6 [n16/test n17/test n18/test]
7 [n19/test n20/test]
8 [n21/test n22/test n23/test]
9 [n24/test n25/test]}
(partition-tests* 10 (map #(symbol (format "n%02d/test" %)) (range 26)))))))

(deftest ^:parallel partition-should-not-split-in-the-middle-of-a-namespace-test
(testing "Partitioning should not split in the middle of a namespace"
(is (= '{0 [a/test-1 a/test-2 a/test-3]
1 [b/test-1]}
(partition-tests* 2 '[a/test-1 a/test-2 a/test-3 b/test-1])))))

(deftest ^:parallel partition-preserve-order-test
(testing "Partitioning should preserve order of namespaces and vars"
(is (= '{0 [b/test-1 b/test-2 b/test-3]
1 [a/test-1 a/test-3 a/test-2]}
(partition-tests* 2 '[b/test-1 b/test-2 b/test-3 a/test-1 a/test-3 a/test-2])))))

(deftest ^:parallel partition-test
(is (= {0 [#'find-tests-test
#'exclude-tags-test]
1 [#'mb.hawk.speak-test/speak-results-test]
2 [#'mb.hawk.parallel-test/ns-parallel-test
#'mb.hawk.parallel-test/var-not-parallel-test]}
(into (sorted-map)
(map (fn [i]
[i (hawk/find-tests-with-options
{:only [`find-tests-test
'mb.hawk.speak-test/speak-results-test
;; this var intentionally comes after a different var in a different
;; namespace to make sure we partition things in a way that groups
;; namespaces together
`exclude-tags-test
'mb.hawk.parallel-test/ns-parallel-test
'mb.hawk.parallel-test/var-not-parallel-test]
:partition/index i
:partition/total 3})]))
(range 3)))))
Loading

0 comments on commit 828d000

Please sign in to comment.