From ef6a42b7bc41f3e247e60927010d8e663882509d Mon Sep 17 00:00:00 2001
From: Matt Brictson <matt@mattbrictson.com>
Date: Sun, 25 Feb 2024 09:43:47 -0800
Subject: [PATCH] Update shard logic to evenly distribute slow tests (#17)

When distributing test files across shards, mighty_test will now ensure
that each shard gets a roughly equal number of slow tests.
---
 lib/mighty_test/cli.rb           |  2 +-
 lib/mighty_test/sharder.rb       | 19 ++++++++++++---
 test/mighty_test/sharder_test.rb | 42 +++++++++++++++++++++++++++++++-
 3 files changed, 57 insertions(+), 6 deletions(-)

diff --git a/lib/mighty_test/cli.rb b/lib/mighty_test/cli.rb
index 4b3f99f..cea588d 100644
--- a/lib/mighty_test/cli.rb
+++ b/lib/mighty_test/cli.rb
@@ -54,7 +54,7 @@ def run_test_by_line_number
     def run_tests_by_path
       test_paths = find_test_paths
       test_paths = excluding_slow_paths(test_paths) unless path_args.any? || ci? || options[:all]
-      test_paths = Sharder.from_argv(options[:shard], env:).shard(test_paths) if options[:shard]
+      test_paths = Sharder.from_argv(options[:shard], env:, file_system:).shard(test_paths) if options[:shard]
 
       run_tests_and_exit!(*test_paths)
     end
diff --git a/lib/mighty_test/sharder.rb b/lib/mighty_test/sharder.rb
index 2981fd2..20a37be 100644
--- a/lib/mighty_test/sharder.rb
+++ b/lib/mighty_test/sharder.rb
@@ -2,19 +2,19 @@ module MightyTest
   class Sharder
     DEFAULT_SEED = 123_456_789
 
-    def self.from_argv(value, env: ENV)
+    def self.from_argv(value, env: ENV, file_system: FileSystem.new)
       index, total = value.to_s.match(%r{\A(\d+)/(\d+)\z})&.captures&.map(&:to_i)
       raise ArgumentError, "shard: value must be in the form INDEX/TOTAL (e.g. 2/8)" if total.nil?
 
       git_sha = env.values_at("GITHUB_SHA", "CIRCLE_SHA1").find { |sha| !sha.to_s.strip.empty? }
       seed = git_sha&.unpack1("l_")
 
-      new(index:, total:, seed:)
+      new(index:, total:, seed:, file_system:)
     end
 
     attr_reader :index, :total, :seed
 
-    def initialize(index:, total:, seed: nil)
+    def initialize(index:, total:, seed: nil, file_system: FileSystem.new)
       raise ArgumentError, "shard: total shards must be a number greater than 0" unless total > 0
 
       valid_group = index > 0 && index <= total
@@ -23,13 +23,24 @@ def initialize(index:, total:, seed: nil)
       @index = index
       @total = total
       @seed = seed || DEFAULT_SEED
+      @file_system = file_system
     end
 
     def shard(*test_paths)
       random = Random.new(seed)
-      shuffled_paths = test_paths.flatten.shuffle(random:)
+
+      # Shuffle slow and normal paths separately so that slow ones get evenly distributed
+      shuffled_paths = test_paths
+        .flatten
+        .partition { |path| !file_system.slow_test_path?(path) }
+        .flat_map { |paths| paths.shuffle(random:) }
+
       slices = shuffled_paths.each_slice(total)
       slices.filter_map { |slice| slice[index - 1] }
     end
+
+    private
+
+    attr_reader :file_system
   end
 end
diff --git a/test/mighty_test/sharder_test.rb b/test/mighty_test/sharder_test.rb
index a4d7bb5..39adc5f 100644
--- a/test/mighty_test/sharder_test.rb
+++ b/test/mighty_test/sharder_test.rb
@@ -55,7 +55,7 @@ def test_for_a_given_seed_it_generates_a_stable_shuffled_result
       assert_equal(%w[f e c], result)
     end
 
-    def test_it_divdes_items_into_roughly_equally_sized_shards
+    def test_it_divides_items_into_roughly_equally_sized_shards
       all = %w[a b c d e f g h i j k l m n o p q r]
       shards = (1..4).map do |index|
         Sharder.new(index:, total: 4).shard(all)
@@ -67,5 +67,45 @@ def test_it_divdes_items_into_roughly_equally_sized_shards
 
       assert_equal all, shards.flatten.sort
     end
+
+    def test_it_evenly_distributes_slow_paths_across_shards
+      all = %w[
+        test/system/login_test.rb
+        test/system/admin_test.rb
+        test/models/post_test.rb
+        test/system/editor_test.rb
+        test/models/user_test.rb
+        test/system/email_test.rb
+        test/models/comment_test.rb
+        test/system/rss_test.rb
+        test/models/category_test.rb
+        test/system/moderation_test.rb
+      ]
+      shards = (1..3).map do |index|
+        Sharder.new(index:, total: 3).shard(all)
+      end
+
+      assert_equal(
+        [
+          %w[
+            test/models/user_test.rb
+            test/models/post_test.rb
+            test/system/login_test.rb
+            test/system/admin_test.rb
+          ],
+          %w[
+            test/models/comment_test.rb
+            test/system/rss_test.rb
+            test/system/moderation_test.rb
+          ],
+          %w[
+            test/models/category_test.rb
+            test/system/email_test.rb
+            test/system/editor_test.rb
+          ]
+        ],
+        shards
+      )
+    end
   end
 end