From b0356b34b5458481c8afaa6b4120be8a61490408 Mon Sep 17 00:00:00 2001
From: Zarif Mahfuz <zarif.mahfuz@shopify.com>
Date: Thu, 27 Mar 2025 10:28:00 -0600
Subject: [PATCH] Specify flaky tests in junit.xml format

---
 ruby/lib/ci/queue/configuration.rb       |  9 ++++++++-
 ruby/test/ci/queue/configuration_test.rb | 15 +++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/ruby/lib/ci/queue/configuration.rb b/ruby/lib/ci/queue/configuration.rb
index 1941b81e..75455437 100644
--- a/ruby/lib/ci/queue/configuration.rb
+++ b/ruby/lib/ci/queue/configuration.rb
@@ -25,7 +25,14 @@ def from_env(env)
 
         def load_flaky_tests(path)
           return [] unless path
-          ::File.readlines(path).map(&:chomp).to_set
+          if ::File.extname(path) == ".xml"
+            require 'rexml/document'
+            REXML::Document.new(::File.read(path)).elements.to_a("//testcase").map do |element|
+              "#{element.attributes['classname']}##{element.attributes['name']}"
+            end.to_set
+          else
+            ::File.readlines(path).map(&:chomp).to_set
+          end
         rescue SystemCallError
           []
         end
diff --git a/ruby/test/ci/queue/configuration_test.rb b/ruby/test/ci/queue/configuration_test.rb
index f79c6aef..18fce529 100644
--- a/ruby/test/ci/queue/configuration_test.rb
+++ b/ruby/test/ci/queue/configuration_test.rb
@@ -90,6 +90,21 @@ def test_parses_file_correctly
 
       flaky_tests = Configuration.load_flaky_tests('/tmp/does-not-exist')
       assert_empty flaky_tests
+
+      Tempfile.open(['flaky_test_file', '.junit.xml']) do |file|
+        file.write(<<~XML)
+          <testsuite name="ATest">
+            <testcase name="test_foo" classname="ATest" />
+            <testcase name="test_bar" classname="ATest" />
+          </testsuite>
+        XML
+        file.close
+
+        flaky_tests = Configuration.load_flaky_tests(file.path)
+        assert_equal 2, flaky_tests.size
+        assert_includes flaky_tests, "ATest#test_foo"
+        assert_includes flaky_tests, "ATest#test_bar"
+      end
     end
 
     def test_queue_init_timeout_unset