Use TestNG retry to reduce flakiness in probabilistic result validation

Author: hantangwangd

presto-main-base/src/test/java/com/facebook/presto/operator/aggregation/noisyaggregation/TestNoisySumGaussianLongAggregation.java

@@ -24,6 +24,8 @@
 import com.facebook.presto.testing.LocalQueryRunner;
 import com.facebook.presto.testing.MaterializedResult;
 import com.facebook.presto.testing.MaterializedRow;
+import com.facebook.presto.util.RetryAnalyzer;
+import com.facebook.presto.util.RetryCount;
 import org.testng.annotations.Test;
 
 import java.util.Arrays;
@@ -422,7 +424,8 @@ public void testNoisySumGaussianLongClippingSomeNoiseScale()
                 expected);
     }
 
-    @Test
+    @Test(retryAnalyzer = RetryAnalyzer.class)
+    @RetryCount(100)
     public void testNoisySumGaussianLongClippingSomeNoiseScaleWithinSomeStd()
     {
         JavaAggregationFunctionImplementation noisySumGaussian = getFunction(BIGINT, DOUBLE, DOUBLE, DOUBLE);

presto-main-base/src/test/java/com/facebook/presto/operator/aggregation/noisyaggregation/TestNoisySumGaussianLongDecimalAggregation.java

@@ -25,6 +25,8 @@
 import com.facebook.presto.testing.LocalQueryRunner;
 import com.facebook.presto.testing.MaterializedResult;
 import com.facebook.presto.testing.MaterializedRow;
+import com.facebook.presto.util.RetryAnalyzer;
+import com.facebook.presto.util.RetryCount;
 import org.testng.annotations.Test;
 
 import java.util.Arrays;
@@ -281,7 +283,8 @@ public void testNoisySumGaussianLongDecimalClippingSomeNoiseScale()
                 expected);
     }
 
-    @Test
+    @Test(retryAnalyzer = RetryAnalyzer.class)
+    @RetryCount(100)
     public void testNoisySumGaussianLongDecimalClippingSomeNoiseScaleWithinSomeStd()
     {
         JavaAggregationFunctionImplementation noisySumGaussian = getFunction(LONG_DECIMAL_TYPE, DOUBLE, DOUBLE, DOUBLE);

presto-main-base/src/test/java/com/facebook/presto/util/RetryAnalyzer.java

@@ -0,0 +1,39 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.util;
+
+import org.testng.IRetryAnalyzer;
+import org.testng.ITestResult;
+
+import java.lang.reflect.Method;
+
+import static com.facebook.presto.util.RetryCount.DEFAULT_RETRY_COUNT;
+
+public class RetryAnalyzer
+        implements IRetryAnalyzer
+{
+    private int retryCount;
+    private int maxRetryCount;
+
+    @Override
+    public boolean retry(ITestResult result)
+    {
+        if (maxRetryCount == 0) {
+            Method method = result.getMethod().getConstructorOrMethod().getMethod();
+            RetryCount annotation = method.getAnnotation(RetryCount.class);
+            maxRetryCount = (annotation != null) ? annotation.value() : DEFAULT_RETRY_COUNT;
+        }
+        return retryCount++ < maxRetryCount;
+    }
+}

presto-main-base/src/test/java/com/facebook/presto/util/RetryCount.java

@@ -0,0 +1,24 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.util;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
+@Retention(RetentionPolicy.RUNTIME)
+public @interface RetryCount
+{
+    int DEFAULT_RETRY_COUNT = 3;
+    int value() default DEFAULT_RETRY_COUNT;
+}