From 79c6a6fa02e58aa9432eb2463d093221903cfead Mon Sep 17 00:00:00 2001
From: Kai Yang <kfstorm@outlook.com>
Date: Mon, 28 Sep 2020 22:16:06 +0800
Subject: [PATCH] [Java] Enable retry in TestNG (#11065)

* Enable retry in Java test

* lint

* update
---
 .../io/ray/test/AnnotationTransformer.java    | 15 +++++++++++++++
 .../main/java/io/ray/test/RetryAnalyzer.java  | 19 +++++++++++++++++++
 java/testng.xml                               |  1 +
 3 files changed, 35 insertions(+)
 create mode 100644 java/test/src/main/java/io/ray/test/AnnotationTransformer.java
 create mode 100644 java/test/src/main/java/io/ray/test/RetryAnalyzer.java

diff --git a/java/test/src/main/java/io/ray/test/AnnotationTransformer.java b/java/test/src/main/java/io/ray/test/AnnotationTransformer.java
new file mode 100644
index 000000000..57e5b35e9
--- /dev/null
+++ b/java/test/src/main/java/io/ray/test/AnnotationTransformer.java
@@ -0,0 +1,15 @@
+package io.ray.test;
+
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Method;
+import org.testng.IAnnotationTransformer;
+import org.testng.annotations.ITestAnnotation;
+
+public class AnnotationTransformer implements IAnnotationTransformer {
+
+  @Override
+  public void transform(ITestAnnotation annotation, Class testClass,
+      Constructor testConstructor, Method testMethod) {
+    annotation.setRetryAnalyzer(RetryAnalyzer.class);
+  }
+}
diff --git a/java/test/src/main/java/io/ray/test/RetryAnalyzer.java b/java/test/src/main/java/io/ray/test/RetryAnalyzer.java
new file mode 100644
index 000000000..5b5b48713
--- /dev/null
+++ b/java/test/src/main/java/io/ray/test/RetryAnalyzer.java
@@ -0,0 +1,19 @@
+package io.ray.test;
+
+import org.testng.IRetryAnalyzer;
+import org.testng.ITestResult;
+
+public class RetryAnalyzer implements IRetryAnalyzer {
+
+  private int counter = 0;
+  private static final int RETRY_LIMIT = 2;
+
+  @Override
+  public boolean retry(ITestResult result) {
+    if (counter < RETRY_LIMIT) {
+      counter++;
+      return true;
+    }
+    return false;
+  }
+}
diff --git a/java/testng.xml b/java/testng.xml
index d81b75252..aae887f2a 100644
--- a/java/testng.xml
+++ b/java/testng.xml
@@ -10,5 +10,6 @@
     <listeners>
         <listener class-name="io.ray.test.RayAlterSuiteListener" />
         <listener class-name="io.ray.test.TestProgressListener" />
+        <listener class-name="io.ray.test.AnnotationTransformer" />
     </listeners>
 </suite>