Use reflection to inject the attempt number

This CL is to address the public static field in RepeatableRunner
for caller to get the current attempt number. We tried to have
a JUnit TestRule to achieve the purpose but it ended up with having
a RuleChain in each class where we already have multiple rules and
need to add the retry rule. This is because we have to make sure
the retry rule is the last one to wrap the test statement so that
the actual retry can include the actions defined in other rules.
Having a rule chain is not scalable and confuses engineer so we
gave it up.

Instead, we decided to expand the current RepeatableRunner to
use reflection to inject the attempt number to the test class.
Doing it this way can reduce the burden from the caller and it also
gets rid of the global state from the previous public static field.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=240789045

javatests/google/registry/webdriver/RepeatableRunner.java

@@ -17,6 +17,10 @@ package google.registry.webdriver;
 import static com.google.common.base.Preconditions.checkState;
 
 import com.google.common.flogger.FluentLogger;
+import java.lang.reflect.Field;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import org.junit.runners.BlockJUnit4ClassRunner;
 import org.junit.runners.model.FrameworkMethod;
 import org.junit.runners.model.InitializationError;
@@ -25,6 +29,21 @@ import org.junit.runners.model.Statement;
 /**
  * A JUnit test runner which can retry each test up to 3 times.
  *
+ * <p>To use this runner, annotate the test class with {@link RepeatableRunner} and define a field
+ * with type of {@link AttemptNumber}:
+ *
+ * <pre>
+ * &#064;RunWith(RepeatableRunner.class)
+ * public class RepeatableRunnerTest {
+ *   private AttemptNumber attemptNumber = new AttemptNumber();
+ *
+ *   &#064;Test
+ *   public void test() {
+ *     print(attemptNumber.get());
+ *   }
+ * }
+ * </pre>
+ *
  * <p>This runner is for our visual regression to prevent flakes during the run. The test is
  * repeated multiple times until it passes and it only fails if it failed 3 times.
  */
@@ -38,51 +57,102 @@ public class RepeatableRunner extends BlockJUnit4ClassRunner {
   private static final int MAX_ATTEMPTS =
       Integer.parseInt(System.getProperty("test.screenshot.maxAttempts", "3"));
 
-  // TODO(b/127984872): Find an elegant way to pass the index of attempt to the test
-  public static int currentAttemptIndex;
+  private final Field attemptNumberField;
+  private AttemptNumber lastAttemptNumber;
 
   /** Constructs a new instance of the default runner */
-  public RepeatableRunner(Class<?> klass) throws InitializationError {
-    super(klass);
+  public RepeatableRunner(Class<?> clazz) throws InitializationError {
+    super(clazz);
+    attemptNumberField = getAttemptNumberField();
+  }
+
+  private Field getAttemptNumberField() {
+    List<Field> attemptNumberFields =
+        Stream.of(getTestClass().getJavaClass().getDeclaredFields())
+            .filter(declaredField -> declaredField.getType().equals(AttemptNumber.class))
+            .collect(Collectors.toList());
+
+    if (attemptNumberFields.size() == 0) {
+      throw new IllegalArgumentException("Missing a field with type of AttemptNumber");
+    } else if (attemptNumberFields.size() > 1) {
+      throw new IllegalArgumentException(
+          "Cannot have more than 1 field with type of AttemptNumber");
+    }
+    Field attemptNumberField = attemptNumberFields.get(0);
+    // It should not matter if that field is set to private
+    attemptNumberField.setAccessible(true);
+    return attemptNumberField;
+  }
+
+  @Override
+  protected Object createTest() throws Exception {
+    Object testObject = super.createTest();
+    // lastAttemptNumber must be null at this moment to indicate that we have
+    // created the RepeatableStatement for the previous AttemptNumber object
+    // or this is the first time we run the test.
+    // If it is not the case, either the tests are run in parallel or the
+    // behavior of BlockJUnit4ClassRunner is changed.
+    checkState(lastAttemptNumber == null);
+    lastAttemptNumber = (AttemptNumber) attemptNumberField.get(testObject);
+    return testObject;
   }
 
   @Override
   protected Statement methodBlock(FrameworkMethod method) {
     Statement statement = super.methodBlock(method);
-    return new RepeatableStatement(statement, method);
+    RepeatableStatement repeatableStatement =
+        new RepeatableStatement(statement, method, lastAttemptNumber);
+    // Explicitly set lastAttemptNumber to null because it should
+    // not be reused accidentally by the next test.
+    lastAttemptNumber = null;
+    return repeatableStatement;
+  }
+
+  /** A simple POJO to store the number of the test attempt. */
+  public static class AttemptNumber {
+    private int attemptNumber;
+
+    /** Returns the number of the test attempt. */
+    public int get() {
+      return attemptNumber;
+    }
+
+    private void set(int attemptNumber) {
+      this.attemptNumber = attemptNumber;
+    }
   }
 
   private static class RepeatableStatement extends Statement {
 
     private Statement statement;
     private FrameworkMethod method;
+    private AttemptNumber attemptNumber;
 
-    public RepeatableStatement(Statement statement, FrameworkMethod method) {
+    public RepeatableStatement(
+        Statement statement, FrameworkMethod method, AttemptNumber attemptNumber) {
       this.statement = statement;
       this.method = method;
+      this.attemptNumber = attemptNumber;
     }
 
     @Override
     public void evaluate() throws Throwable {
       checkState(MAX_ATTEMPTS > 0);
       int numSuccess = 0, numFailure = 0;
-      Throwable firstException = null;
+      Throwable lastException = null;
       for (int attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) {
-        currentAttemptIndex = attempt;
+        attemptNumber.set(attempt);
         try {
           statement.evaluate();
           numSuccess++;
-          if (RUN_ALL_ATTEMPTS) {
-            logger.atInfo().log(
-                "[%s] Attempt %d of %d succeeded!\n", method.getName(), attempt, MAX_ATTEMPTS);
-            continue;
+          logger.atInfo().log(
+              "[%s] Attempt %d of %d succeeded!\n", method.getName(), attempt, MAX_ATTEMPTS);
+          if (!RUN_ALL_ATTEMPTS) {
+            return;
           }
-          return;
         } catch (Throwable e) {
           numFailure++;
-          if (firstException == null) {
-            firstException = e;
-          }
+          lastException = e;
           logger.atWarning().log(
               "[%s] Attempt %d of %d failed!\n", method.getName(), attempt, MAX_ATTEMPTS);
         }
@@ -93,7 +163,7 @@ public class RepeatableRunner extends BlockJUnit4ClassRunner {
       if (numSuccess == 0) {
         logger.atSevere().log(
             "[%s] didn't pass after all %d attempts failed!\n", method.getName(), MAX_ATTEMPTS);
-        throw firstException;
+        throw lastException;
       }
     }
   }