[FLINK-20723][tests] Allow retries to be defined per class

flink-test-utils-parent/flink-test-utils-junit/src/main/java/org/apache/flink/testutils/junit/RetryOnException.java

@@ -18,30 +18,38 @@
 
 package org.apache.flink.testutils.junit;
 
+import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
 /**
  * Annotation to use with {@link org.apache.flink.testutils.junit.RetryRule}.
  *
- * <p>Add the {@link org.apache.flink.testutils.junit.RetryRule} to your test and annotate tests
- * with {@link RetryOnException}.
+ * <p>Add the {@link org.apache.flink.testutils.junit.RetryRule} to your test class and annotate the
+ * class and/or tests with {@link RetryOnException}.
  *
  * <pre>
+ * {@literal @}RetryOnException(times=1, exception=IOException.class)
  * public class YourTest {
  *
  *     {@literal @}Rule
  *     public RetryRule retryRule = new RetryRule();
  *
  *     {@literal @}Test
- *     {@literal @}RetryOnException(times=1, exception=IOException.class)
  *     public void yourTest() throws Exception {
  *         // This will be retried 1 time (total runs 2) before failing the test.
  *         throw new IOException("Failing test");
  *     }
  *
  *     {@literal @}Test
+ *     {@literal @}RetryOnException(times=2, exception=IOException.class)
+ *     public void yourTest() throws Exception {
+ *         // This will be retried 2 times (total runs 3) before failing the test.
+ *         throw new IOException("Failing test");
+ *     }
+ *
+ *     {@literal @}Test
  *     {@literal @}RetryOnException(times=1, exception=IOException.class)
  *     public void yourTest() throws Exception {
  *         // This will not be retried, because it throws the wrong exception
@@ -51,7 +59,7 @@
  * </pre>
  */
 @Retention(RetentionPolicy.RUNTIME)
-@Target(java.lang.annotation.ElementType.METHOD)
+@Target({java.lang.annotation.ElementType.METHOD, ElementType.TYPE})
 public @interface RetryOnException {
 
     int times();

flink-test-utils-parent/flink-test-utils-junit/src/main/java/org/apache/flink/testutils/junit/RetryOnFailure.java

@@ -18,32 +18,41 @@
 
 package org.apache.flink.testutils.junit;
 
+import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
 /**
  * Annotation to use with {@link RetryRule}.
  *
- * <p>Add the {@link RetryRule} to your test and annotate tests with {@link RetryOnFailure}.
+ * <p>Add the {@link RetryRule} to your test class and annotate the class and/or tests with {@link
+ * RetryOnFailure}.
  *
  * <pre>
+ * {@literal @}RetryOnFailure(times=1)
  * public class YourTest {
  *
  *     {@literal @}Rule
  *     public RetryRule retryRule = new RetryRule();
  *
  *     {@literal @}Test
- *     {@literal @}RetryOnFailure(times=1)
  *     public void yourTest() {
  *         // This will be retried 1 time (total runs 2) before failing the test.
  *         throw new Exception("Failing test");
  *     }
+ *
+ *     {@literal @}Test
+ *     {@literal @}RetryOnFailure(times=2)
+ *     public void yourTest() {
+ *         // This will be retried 2 time (total runs 3) before failing the test.
+ *         throw new Exception("Failing test");
+ *     }
  * }
  * </pre>
  */
 @Retention(RetentionPolicy.RUNTIME)
-@Target({java.lang.annotation.ElementType.METHOD})
+@Target({java.lang.annotation.ElementType.METHOD, ElementType.TYPE})
 public @interface RetryOnFailure {
     int times();
 }

flink-test-utils-parent/flink-test-utils-junit/src/main/java/org/apache/flink/testutils/junit/RetryRule.java

@@ -25,23 +25,34 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.annotation.Nullable;
+
 /**
  * A rule to retry failed tests for a fixed number of times.
  *
- * <p>Add the {@link RetryRule} to your test and annotate tests with {@link RetryOnFailure}.
+ * <p>Add the {@link RetryRule} to your test class and annotate the class and/or tests with either
+ * {@link RetryOnFailure} or {@link RetryOnException}. If both the class and test are annotated,
+ * then only the latter annotation is taken into account.
  *
  * <pre>
+ * {@literal @}RetryOnFailure(times=1)
  * public class YourTest {
  *
  *     {@literal @}Rule
  *     public RetryRule retryRule = new RetryRule();
  *
  *     {@literal @}Test
- *     {@literal @}RetryOnFailure(times=1)
  *     public void yourTest() {
  *         // This will be retried 1 time (total runs 2) before failing the test.
  *         throw new Exception("Failing test");
  *     }
+ *
+ *     {@literal @}Test
+ *     {@literal @}RetryOnFailure(times=2)
+ *     public void yourTest() {
+ *         // This will be retried 2 time (total runs 3) before failing the test.
+ *         throw new Exception("Failing test");
+ *     }
  * }
  * </pre>
  */
@@ -54,6 +65,12 @@ public Statement apply(Statement statement, Description description) {
         RetryOnFailure retryOnFailure = description.getAnnotation(RetryOnFailure.class);
         RetryOnException retryOnException = description.getAnnotation(RetryOnException.class);
 
+        if (retryOnFailure == null && retryOnException == null) {
+            // if nothing is specified on the test method, fall back to annotations on the class
+            retryOnFailure = description.getTestClass().getAnnotation(RetryOnFailure.class);
+            retryOnException = description.getTestClass().getAnnotation(RetryOnException.class);
+        }
+
         // sanity check that we don't use both annotations
         if (retryOnFailure != null && retryOnException != null) {
             throw new IllegalArgumentException(

flink-test-utils-parent/flink-test-utils-junit/src/test/java/org/apache/flink/testutils/junit/RetryRuleTest.java

@@ -30,10 +30,11 @@
 
 /** Tests for the {@link RetryRule}. */
 public class RetryRuleTest extends TestLogger {
+    private static final RetryRule RETRY_RULE = new RetryRule();
 
     @Test
     public void testExpectedExceptionIgnored() throws Throwable {
-        final RetryRule retryRule = new RetryRule();
+        final int numEvaluationsToFail = 1;
 
         final Description testDescription =
                 Description.createTestDescription(
@@ -43,15 +44,15 @@ public void testExpectedExceptionIgnored() throws Throwable {
                                 .getMethod("test")
                                 .getAnnotations());
 
-        final TestStatement statement = new TestStatement(1);
+        final TestStatement statement = new TestStatement(numEvaluationsToFail);
 
         try {
-            retryRule.apply(statement, testDescription).evaluate();
+            RETRY_RULE.apply(statement, testDescription).evaluate();
             Assert.fail("Should have failed.");
         } catch (RuntimeException expected) {
         }
 
-        assertThat(statement.getNumEvaluations(), is(1));
+        assertThat(statement.getNumEvaluations(), is(numEvaluationsToFail));
     }
 
     @Ignore // we don't want to actually this run as a test
@@ -61,6 +62,82 @@ private static class TestClassWithTestExpectingRuntimeException {
         public void test() {}
     }
 
+    @Test
+    public void testNoAnnotationResultsInZeroRetries() throws Throwable {
+        final int numEvaluationsToFail = 1;
+
+        final Description testDescription =
+                Description.createTestDescription(
+                        TestClassWithoutAnnotation.class,
+                        "test",
+                        TestClassWithAnnotation.class.getMethod("test").getAnnotations());
+
+        final TestStatement statement = new TestStatement(numEvaluationsToFail);
+
+        try {
+            RETRY_RULE.apply(statement, testDescription).evaluate();
+            Assert.fail("Should have failed.");
+        } catch (RuntimeException expected) {
+        }
+
+        assertThat(statement.getNumEvaluations(), is(numEvaluationsToFail));
+    }
+
+    @Ignore // we don't want to actually this run as a test
+    private static class TestClassWithoutAnnotation {
+        @Test
+        public void test() {}
+    }
+
+    @Test
+    public void testAnnotationOnClassUsedAsFallback() throws Throwable {
+        final int numEvaluationsToFail = 1;
+
+        final Description testDescription =
+                Description.createTestDescription(
+                        TestClassWithAnnotation.class,
+                        "test",
+                        TestClassWithAnnotation.class.getMethod("test").getAnnotations());
+
+        final TestStatement statement = new TestStatement(numEvaluationsToFail);
+
+        RETRY_RULE.apply(statement, testDescription).evaluate();
+
+        assertThat(statement.getNumEvaluations(), is(numEvaluationsToFail + 1));
+    }
+
+    @Ignore // we don't want to actually this run as a test
+    @RetryOnFailure(times = 1)
+    private static class TestClassWithAnnotation {
+        @Test
+        public void test() {}
+    }
+
+    @Test
+    public void testAnnotationOnMethodTakesPrecedence() throws Throwable {
+        final int numEvaluationsToFail = 2;
+
+        final Description testDescription =
+                Description.createTestDescription(
+                        TestClassWithAnnotationOnMethod.class,
+                        "test",
+                        TestClassWithAnnotationOnMethod.class.getMethod("test").getAnnotations());
+
+        final TestStatement statement = new TestStatement(numEvaluationsToFail);
+
+        RETRY_RULE.apply(statement, testDescription).evaluate();
+
+        assertThat(statement.getNumEvaluations(), is(numEvaluationsToFail + 1));
+    }
+
+    @Ignore // we don't want to actually this run as a test
+    @RetryOnFailure(times = 1)
+    private static class TestClassWithAnnotationOnMethod {
+        @RetryOnFailure(times = 2)
+        @Test
+        public void test() {}
+    }
+
     private static class TestStatement extends Statement {
         private final int numEvaluationsToFail;