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/ActualScreenshot.java

@@ -31,14 +31,16 @@ public class ActualScreenshot {
   public static final String IMAGE_FORMAT = "png";
   private String imageKey;
   private BufferedImage bufferedImage;
+  private int attempt;
 
-  private ActualScreenshot(String imageKey, BufferedImage bufferedImage) {
+  private ActualScreenshot(String imageKey, BufferedImage bufferedImage, int attempt) {
     this.imageKey = imageKey;
     this.bufferedImage = bufferedImage;
+    this.attempt = attempt;
   }
 
   /** Creates an ActualScreenshot from the given image format and byte array. */
-  public static ActualScreenshot create(String imageKey, byte[] imageBytes) {
+  public static ActualScreenshot create(String imageKey, int attempt, byte[] imageBytes) {
     checkNotNull(imageKey);
     checkNotNull(imageBytes);
     byte[] imageBytesClone = Arrays.copyOf(imageBytes, imageBytes.length);
@@ -47,7 +49,7 @@ public class ActualScreenshot {
     try {
       imageReader.setInput(checkNotNull(ImageIO.createImageInputStream(imageInputStream)));
       BufferedImage bufferedImage = checkNotNull(imageReader.read(0));
-      return new ActualScreenshot(imageKey, bufferedImage);
+      return new ActualScreenshot(imageKey, bufferedImage, attempt);
     } catch (IOException e) {
       throw new UncheckedIOException(e);
     }
@@ -55,7 +57,7 @@ public class ActualScreenshot {
 
   /** {@link BufferedImage#getSubimage(int, int, int, int)} */
   public ActualScreenshot getSubimage(int x, int y, int w, int h) {
-    return new ActualScreenshot(imageKey, bufferedImage.getSubimage(x, y, w, h));
+    return new ActualScreenshot(imageKey, bufferedImage.getSubimage(x, y, w, h), attempt);
   }
 
   /** {@link BufferedImage#getWidth()} */
@@ -82,9 +84,9 @@ public class ActualScreenshot {
     }
   }
 
-  /** Returns the imageKey of the screenshot. */
+  /** Returns the attempt number of the test. */
-  public String getImageKey() {
+  public int getAttempt() {
-    return imageKey;
+    return attempt;
   }
 
   /** Returns the concat of imageKey and imageFormat. */

javatests/google/registry/webdriver/ChromeWebDriverPlusScreenDiffer.java

@@ -16,7 +16,6 @@ package google.registry.webdriver;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
-import static google.registry.webdriver.RepeatableRunner.currentAttemptIndex;
 import static java.lang.Math.abs;
 import static java.util.stream.Collectors.joining;
 
@@ -123,9 +122,9 @@ class ChromeWebDriverPlusScreenDiffer implements WebDriverPlusScreenDiffer {
   }
 
   @Override
-  public void diffElement(String imageKey, WebElement element) {
+  public void diffElement(WebElement element, String imageKey, int attempt) {
     ActualScreenshot elementImage =
-        takeScreenshot(imageKey)
+        ActualScreenshot.create(imageKey, attempt, takeScreenshot())
             .getSubimage(
                 element.getLocation().getX(),
                 element.getLocation().getY(),
@@ -135,8 +134,8 @@ class ChromeWebDriverPlusScreenDiffer implements WebDriverPlusScreenDiffer {
   }
 
   @Override
-  public void diffPage(String imageKey) {
+  public void diffPage(String imageKey, int attempt) {
-    actualScreenshots.add(takeScreenshot(imageKey));
+    actualScreenshots.add(ActualScreenshot.create(imageKey, attempt, takeScreenshot()));
   }
 
   @Override
@@ -190,10 +189,10 @@ class ChromeWebDriverPlusScreenDiffer implements WebDriverPlusScreenDiffer {
     }
   }
 
-  private ActualScreenshot takeScreenshot(String imageKey) {
+  private byte[] takeScreenshot() {
     checkArgument(webDriver instanceof TakesScreenshot);
     TakesScreenshot takesScreenshot = (TakesScreenshot) webDriver;
-    return ActualScreenshot.create(imageKey, takesScreenshot.getScreenshotAs(OutputType.BYTES));
+    return takesScreenshot.getScreenshotAs(OutputType.BYTES);
   }
 
   private ComparisonResult compareScreenshots(ActualScreenshot screenshot) {
@@ -246,7 +245,7 @@ class ChromeWebDriverPlusScreenDiffer implements WebDriverPlusScreenDiffer {
     File thisScreenshotDir =
         Paths.get(
                 screenshotDir,
-                "attempt_" + currentAttemptIndex,
+                "attempt_" + result.actualScreenshot().getAttempt(),
                 result.isConsideredSimilar() ? "similar" : "different")
             .toFile();
     thisScreenshotDir.mkdirs();

javatests/google/registry/webdriver/OteSetupConsoleScreenshotTest.java

@@ -21,6 +21,7 @@ import com.googlecode.objectify.ObjectifyFilter;
 import google.registry.model.ofy.OfyFilter;
 import google.registry.module.frontend.FrontendServlet;
 import google.registry.server.RegistryTestServer;
+import google.registry.webdriver.RepeatableRunner.AttemptNumber;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -40,7 +41,8 @@ public class OteSetupConsoleScreenshotTest {
           .setEmail("Marla.Singer@google.com")
           .build();
 
-  @Rule public final WebDriverRule driver = new WebDriverRule();
+  private final AttemptNumber attemptNumber = new AttemptNumber();
+  @Rule public final WebDriverRule driver = new WebDriverRule(attemptNumber);
 
   @Test
   public void get_owner_fails() throws Throwable {

javatests/google/registry/webdriver/RegistrarConsoleScreenshotTest.java

@@ -27,6 +27,7 @@ import google.registry.model.registrar.Registrar.State;
 import google.registry.module.frontend.FrontendServlet;
 import google.registry.server.RegistryTestServer;
 import google.registry.testing.CertificateSamples;
+import google.registry.webdriver.RepeatableRunner.AttemptNumber;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -50,7 +51,8 @@ public class RegistrarConsoleScreenshotTest {
           .setEmail("Marla.Singer@google.com")
           .build();
 
-  @Rule public final WebDriverRule driver = new WebDriverRule();
+  private final AttemptNumber attemptNumber = new AttemptNumber();
+  @Rule public final WebDriverRule driver = new WebDriverRule(attemptNumber);
 
   @Test
   public void index_owner() throws Throwable {

javatests/google/registry/webdriver/RegistrarConsoleWebTest.java

@@ -29,15 +29,15 @@ import google.registry.model.registrar.RegistrarContact;
 import google.registry.module.frontend.FrontendServlet;
 import google.registry.server.RegistryTestServer;
 import google.registry.testing.AppEngineRule;
+import google.registry.webdriver.RepeatableRunner.AttemptNumber;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.Timeout;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 import org.openqa.selenium.By;
 
 /** WebDriver tests for Registrar Console UI. */
-@RunWith(JUnit4.class)
+@RunWith(RepeatableRunner.class)
 public class RegistrarConsoleWebTest {
 
   @Rule
@@ -59,7 +59,8 @@ public class RegistrarConsoleWebTest {
           .setEmail("Marla.Singer@google.com")
           .build();
 
-  @Rule public final WebDriverRule driver = new WebDriverRule();
+  private final AttemptNumber attemptNumber = new AttemptNumber();
+  @Rule public final WebDriverRule driver = new WebDriverRule(attemptNumber);
 
 
   @Rule public final Timeout deathClock = new Timeout(60000);

javatests/google/registry/webdriver/RegistrarCreateConsoleScreenshotTest.java

@@ -21,6 +21,7 @@ import com.googlecode.objectify.ObjectifyFilter;
 import google.registry.model.ofy.OfyFilter;
 import google.registry.module.frontend.FrontendServlet;
 import google.registry.server.RegistryTestServer;
+import google.registry.webdriver.RepeatableRunner.AttemptNumber;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -40,7 +41,8 @@ public class RegistrarCreateConsoleScreenshotTest {
           .setEmail("Marla.Singer@google.com")
           .build();
 
-  @Rule public final WebDriverRule driver = new WebDriverRule();
+  private final AttemptNumber attemptNumber = new AttemptNumber();
+  @Rule public final WebDriverRule driver = new WebDriverRule(attemptNumber);
 
   @Test
   public void get_owner_fails() throws Throwable {

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
+  private final Field attemptNumberField;
-  public static int currentAttemptIndex;
+  private AttemptNumber lastAttemptNumber;
 
   /** Constructs a new instance of the default runner */
-  public RepeatableRunner(Class<?> klass) throws InitializationError {
+  public RepeatableRunner(Class<?> clazz) throws InitializationError {
-    super(klass);
+    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(
-            logger.atInfo().log(
+              "[%s] Attempt %d of %d succeeded!\n", method.getName(), attempt, MAX_ATTEMPTS);
-                "[%s] Attempt %d of %d succeeded!\n", method.getName(), attempt, MAX_ATTEMPTS);
+          if (!RUN_ALL_ATTEMPTS) {
-            continue;
+            return;
           }
-          return;
         } catch (Throwable e) {
           numFailure++;
-          if (firstException == null) {
+          lastException = e;
-            firstException = 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;
       }
     }
   }

javatests/google/registry/webdriver/WebDriverPlusScreenDiffer.java

@@ -34,11 +34,12 @@ interface WebDriverPlusScreenDiffer {
    * other screenshot matches can be executed in the same function - allowing you to approve /
    * reject all at once.
    *
+   * @param element the element on the page to be compared
    * @param imageKey a unique name such that by prepending the calling class name and method name in
    *     the format of ClassName_MethodName_<imageKey> will uniquely identify golden image.
-   * @param element the element on the page to be compared
+   * @param attempt the attempt number of the test
    */
-  void diffElement(String imageKey, WebElement element);
+  void diffElement(WebElement element, String imageKey, int attempt);
 
   /**
    * Checks that the screenshot matches the golden image by pixel comparison. {@link
@@ -50,8 +51,9 @@ interface WebDriverPlusScreenDiffer {
    *
    * @param imageKey a unique name such that by prepending the calling class name and method name in
    *     the format of ClassName_MethodName_<imageKey> will uniquely identify golden image.
+   * @param attempt the attempt number of the test
    */
-  void diffPage(String imageKey);
+  void diffPage(String imageKey, int attempt);
 
   /** Asserts that all diffs up to this point have PASSED. */
   void verifyAndQuit();

javatests/google/registry/webdriver/WebDriverRule.java

@@ -15,10 +15,11 @@
 package google.registry.webdriver;
 
 import static com.google.common.io.Resources.getResource;
+import static com.google.common.base.Preconditions.checkNotNull;
 import static java.util.stream.Collectors.joining;
 import static org.apache.commons.text.StringEscapeUtils.escapeEcmaScript;
 
-import com.google.common.base.Preconditions;
+import google.registry.webdriver.RepeatableRunner.AttemptNumber;
 import java.net.URL;
 import java.util.List;
 import java.util.Map;
@@ -68,6 +69,7 @@ public final class WebDriverRule extends ExternalResource
   private static final String GOLDENS_PATH =
       getResource(WebDriverRule.class, "goldens/chrome-linux").getFile();
 
+  private AttemptNumber attemptNumber;
   private WebDriver driver;
   private WebDriverPlusScreenDiffer webDriverPlusScreenDiffer;
 
@@ -75,6 +77,11 @@ public final class WebDriverRule extends ExternalResource
   // starts. Will be added a user-given imageKey as a suffix, and of course a '.png' at the end.
   private String imageNamePrefix = null;
 
+  /** Constructs a {@link WebDriverRule} instance. */
+  public WebDriverRule(AttemptNumber attemptNumber) {
+    this.attemptNumber = attemptNumber;
+  }
+
   @Override
   public Statement apply(Statement base, Description description) {
     if (imageNamePrefix == null) {
@@ -163,7 +170,7 @@ public final class WebDriverRule extends ExternalResource
    * @param element the element on the page to be compared
    */
   public void diffElement(String imageKey, WebElement element) {
-    webDriverPlusScreenDiffer.diffElement(getUniqueName(imageKey), element);
+    webDriverPlusScreenDiffer.diffElement(element, getUniqueName(imageKey), attemptNumber.get());
   }
 
   /**
@@ -192,7 +199,7 @@ public final class WebDriverRule extends ExternalResource
    *     the format of ClassName_MethodName_<imageKey> will uniquely identify golden image.
    */
   public void diffPage(String imageKey) {
-    webDriverPlusScreenDiffer.diffPage(getUniqueName(imageKey));
+    webDriverPlusScreenDiffer.diffPage(getUniqueName(imageKey), attemptNumber.get());
   }
 
   @Override
@@ -291,7 +298,7 @@ public final class WebDriverRule extends ExternalResource
   }
 
   private String getUniqueName(String imageKey) {
-    Preconditions.checkNotNull(imageNamePrefix);
+    checkNotNull(imageNamePrefix);
     return imageNamePrefix + "_" + imageKey;
   }
 }