From 30ba9fb92e615209d310335e55c84d6ef40c2b0f Mon Sep 17 00:00:00 2001
From: Aiden Storey <git@storey.dev>
Date: Wed, 29 Apr 2026 14:52:45 -0400
Subject: [PATCH] Add test id normalizer to configuration

---
 ruby/lib/ci/queue/configuration.rb       | 16 +++++-
 ruby/test/ci/queue/configuration_test.rb | 70 ++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/ruby/lib/ci/queue/configuration.rb b/ruby/lib/ci/queue/configuration.rb
index d097f507..6dfe3a86 100644
--- a/ruby/lib/ci/queue/configuration.rb
+++ b/ruby/lib/ci/queue/configuration.rb
@@ -172,7 +172,21 @@ def max_duration=(duration)
       end
 
       def flaky?(test)
-        @flaky_tests.include?(test.id)
+        lazy_normalized_flaky_tests.include?(normalize_test_id(test.id))
+      end
+
+      def normalize_test_id(id)
+        id
+      end
+
+      def lazy_normalized_flaky_tests
+        @normalized_flaky_tests ||= @flaky_tests.to_set { |test_id| normalize_test_id(test_id) }
+      end
+
+      def test_id_normalizer=(normalizer)
+        normalizer = normalizer.to_proc if Method === normalizer
+        define_singleton_method(:normalize_test_id, normalizer)
+        @normalized_flaky_tests = nil
       end
 
       def seed
diff --git a/ruby/test/ci/queue/configuration_test.rb b/ruby/test/ci/queue/configuration_test.rb
index a6728577..2855e99d 100644
--- a/ruby/test/ci/queue/configuration_test.rb
+++ b/ruby/test/ci/queue/configuration_test.rb
@@ -200,6 +200,76 @@ def test_new_lazy_load_test_helpers_env
       assert_equal ["test/test_helper.rb", "test/support/helper.rb"], config.lazy_load_test_helper_paths
     end
 
+    def test_flaky_without_normalizer
+      config = Configuration.new(flaky_tests: Set["ATest#test_foo", "BTest#test_bar"])
+
+      test = Struct.new(:id).new("ATest#test_foo")
+      assert config.flaky?(test)
+
+      test = Struct.new(:id).new("CTest#test_baz")
+      refute config.flaky?(test)
+    end
+
+    def test_flaky_normalized_lookup_uses_set
+      config = Configuration.new(flaky_tests: Set["ATest#test_foo"])
+      assert_kind_of Set, config.lazy_normalized_flaky_tests
+    end
+
+    def test_flaky_with_normalizer
+      config = Configuration.new(flaky_tests: Set["./test/my_test.rb:ATest#test_foo", "./test/my_test.rb:BTest#test_bar"])
+      config.test_id_normalizer = Module.new {
+        define_method(:normalize) { |id| id.sub(%r{^\./}, "") }
+      }.instance_method(:normalize)
+
+      test = Struct.new(:id).new("test/my_test.rb:ATest#test_foo")
+      assert config.flaky?(test)
+
+      test = Struct.new(:id).new("test/my_test.rb:CTest#test_baz")
+      refute config.flaky?(test)
+    end
+
+    def test_flaky_with_method_normalizer
+      config = Configuration.new(flaky_tests: Set["./test/my_test.rb:ATest#test_foo", "./test/my_test.rb:BTest#test_bar"])
+      normalizer_mod = Module.new do
+        def self.normalize(id)
+          id.sub(%r{^\./}, "")
+        end
+      end
+      config.test_id_normalizer = normalizer_mod.method(:normalize)
+
+      test = Struct.new(:id).new("test/my_test.rb:ATest#test_foo")
+      assert config.flaky?(test)
+
+      test = Struct.new(:id).new("test/my_test.rb:CTest#test_baz")
+      refute config.flaky?(test)
+    end
+
+    def test_flaky_normalizer_applied_to_both_sides
+      config = Configuration.new(flaky_tests: Set["  ATest#test_foo  ", "BTest#test_bar"])
+      config.test_id_normalizer = Module.new {
+        define_method(:normalize) { |id| id.strip }
+      }.instance_method(:normalize)
+
+      test = Struct.new(:id).new("ATest#test_foo")
+      assert config.flaky?(test)
+
+      test = Struct.new(:id).new("  ATest#test_foo  ")
+      assert config.flaky?(test)
+    end
+
+    def test_flaky_normalizer_invalidates_cache
+      config = Configuration.new(flaky_tests: Set["./test/my_test.rb:ATest#test_foo"])
+      test = Struct.new(:id).new("test/my_test.rb:ATest#test_foo")
+
+      refute config.flaky?(test)
+
+      config.test_id_normalizer = Module.new {
+        define_method(:normalize) { |id| id.sub(%r{^\./}, "") }
+      }.instance_method(:normalize)
+
+      assert config.flaky?(test)
+    end
+
     def test_heartbeat_max_test_duration_defaults
       # defaults to timeout*10 when heartbeat is enabled
       config = Configuration.new(timeout: 5, max_missed_heartbeat_seconds: 1)