/*
 * 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 io.prestosql.testng.services;

import com.google.common.annotations.VisibleForTesting;
import io.airlift.log.Logger;
import org.testng.IRetryAnalyzer;
import org.testng.ITestNGMethod;
import org.testng.ITestResult;

import javax.annotation.concurrent.GuardedBy;

import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Map;
import java.util.regex.Pattern;
import java.util.stream.Stream;

import static com.google.common.base.Throwables.getStackTraceAsString;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static java.lang.String.format;

public class FlakyTestRetryAnalyzer
        implements IRetryAnalyzer
{
    private static final Logger log = Logger.get(FlakyTestRetryAnalyzer.class);

    // This property exists so that flaky tests are retried on CI only by default but tests of retrying pass locally as well.
    // TODO replace pom.xml property with explicit invocation of a testng runner (test suite with a test) and amend the retryer behavior on that level
    private static final String ENABLE_PROPERTY = "io.prestosql.testng.services.FlakyTestRetryAnalyzer.enabled";

    @VisibleForTesting
    static final int ALLOWED_RETRIES_COUNT = 2;

    @GuardedBy("this")
    private final Map<String, Long> retryCounter = new HashMap<>();

    @Override
    public boolean retry(ITestResult result)
    {
        if (result.isSuccess()) {
            return false;
        }
        if (!isEnabled()) {
            return false;
        }
        Method javaMethod = result.getMethod().getConstructorOrMethod().getMethod();
        if (javaMethod == null) {
            return false;
        }
        Flaky annotation = javaMethod.getAnnotation(Flaky.class);
        if (annotation == null) {
            return false;
        }
        if (result.getThrowable() == null) {
            return false;
        }
        String stackTrace = getStackTraceAsString(result.getThrowable());
        if (!Pattern.compile(annotation.match()).matcher(stackTrace).find()) {
            return false;
        }

        long retryCount;
        ITestNGMethod method = result.getMethod();
        synchronized (this) {
            String name = getName(method, result.getParameters());
            retryCount = retryCounter.getOrDefault(name, 0L);
            retryCount++;
            if (retryCount > ALLOWED_RETRIES_COUNT) {
                return false;
            }
            retryCounter.put(name, retryCount);
        }
        log.warn(
                result.getThrowable(),
                "Test %s::%s attempt %s failed, retrying...,",
                result.getTestClass().getName(),
                method.getMethodName(),
                retryCount);
        return true;
    }

    private static boolean isEnabled()
    {
        if (System.getProperty(ENABLE_PROPERTY) != null) {
            return Boolean.getBoolean(ENABLE_PROPERTY);
        }

        // Enable retry on CI by default
        return System.getenv("CONTINUOUS_INTEGRATION") != null;
    }

    private static String getName(ITestNGMethod method, Object[] parameters)
    {
        String actualTestClass = method.getTestClass().getName();
        if (parameters.length != 0) {
            return format(
                    "%s::%s(%s)",
                    actualTestClass,
                    method.getMethodName(),
                    String.join(",", Stream.of(parameters).map(Object::toString).collect(toImmutableList())));
        }
        return format("%s::%s", actualTestClass, method.getMethodName());
    }
}

@@ -2021,24 +2021,95 @@
 nabled()) %7B%0A
+            log.info(%22not retrying; FlakyTestRetryAnalyzer disabled%22);%0A
             
@@ -2114,32 +2114,32 @@
   return false;%0A
-
         %7D%0A      
@@ -2243,32 +2243,94 @@
 thod == null) %7B%0A
+            log.info(%22not retrying; cannot get java method%22);%0A
             retu
@@ -2441,32 +2441,101 @@
 tion == null) %7B%0A
+            log.info(%22not retrying; @Flaky annotation not present%22);%0A
             retu
@@ -2591,32 +2591,103 @@
 le() == null) %7B%0A
+            log.info(%22not retrying; throwable not present in result%22);%0A
             retu
@@ -2776,24 +2776,24 @@
 rowable());%0A
-
         if (
@@ -2855,24 +2855,139 @@
 ).find()) %7B%0A
+            log.warn(%22not retrying; stacktrace '%25s' does not match pattern '%25s'%22, stackTrace, annotation.match());%0A
             

Add logging to FlakyTestRetryAnalyzer