/*
 * 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.Boolean.parseBoolean;
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 ENABLED_SYSTEM_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) {
            log.info("not retrying; cannot get java method");
            return false;
        }
        Flaky annotation = javaMethod.getAnnotation(Flaky.class);
        if (annotation == null) {
            log.info("not retrying; @Flaky annotation not present");
            return false;
        }
        if (result.getThrowable() == null) {
            log.info("not retrying; throwable not present in result");
            return false;
        }
        String stackTrace = getStackTraceAsString(result.getThrowable());
        if (!Pattern.compile(annotation.match()).matcher(stackTrace).find()) {
            log.warn("not retrying; stacktrace '%s' does not match pattern '%s'", stackTrace, annotation.match());
            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()
    {
        String enabledSystemPropertyValue = System.getProperty(ENABLED_SYSTEM_PROPERTY);
        if (parseBoolean(enabledSystemPropertyValue)) {
            return true;
        }
        // Enable retry on CI by default
        if (System.getenv("CONTINUOUS_INTEGRATION") != null) {
            return true;
        }
        log.info(
                "not retrying; FlakyTestRetryAnalyzer not enabled because: " +
                        "CONTINUOUS_INTEGRATION environment is not detected, " +
                        "system property '%s' is not set to 'true' (but to '%s')",
                ENABLED_SYSTEM_PROPERTY,
                enabledSystemPropertyValue);
        return false;
    }

    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());
    }
}

@@ -903,24 +903,51 @@
 a.util.Map;%0A
+import java.util.Optional;%0A
 import java.
@@ -966,16 +966,16 @@
 attern;%0A
-
 import j
@@ -1145,54 +1145,8 @@
 st;%0A
-import static java.lang.Boolean.parseBoolean;%0A
 impo
@@ -2045,27 +2045,632 @@
     
-if (!isEnabled()) %7B
+Optional%3CString%3E enabledSystemPropertyValue = Optional.ofNullable(System.getProperty(ENABLED_SYSTEM_PROPERTY));%0A        if (!enabledSystemPropertyValue.map(Boolean::parseBoolean)%0A                .orElse(System.getenv(%22CONTINUOUS_INTEGRATION%22) != null)) %7B%0A            log.info(%0A                    %22FlakyTestRetryAnalyzer not enabled: %22 +%0A                            %22CONTINUOUS_INTEGRATION environment is not detected or %22 +%0A                            %22system property '%25s' is not set to 'true' (actual: %25s)%22,%0A                    ENABLED_SYSTEM_PROPERTY,%0A                    enabledSystemPropertyValue.orElse(%22%3Cnot set%3E%22));
 %0A   
@@ -4229,32 +4229,32 @@
    retryCount);%0A
+
         return t
@@ -4269,748 +4269,8 @@
  %7D%0A%0A
-    private static boolean isEnabled()%0A    %7B%0A        String enabledSystemPropertyValue = System.getProperty(ENABLED_SYSTEM_PROPERTY);%0A        if (parseBoolean(enabledSystemPropertyValue)) %7B%0A            return true;%0A        %7D%0A        // Enable retry on CI by default%0A        if (System.getenv(%22CONTINUOUS_INTEGRATION%22) != null) %7B%0A            return true;%0A        %7D%0A        log.info(%0A                %22not retrying; FlakyTestRetryAnalyzer not enabled because: %22 +%0A                        %22CONTINUOUS_INTEGRATION environment is not detected, %22 +%0A                        %22system property '%25s' is not set to 'true' (but to '%25s')%22,%0A                ENABLED_SYSTEM_PROPERTY,%0A                enabledSystemPropertyValue);%0A        return false;%0A    %7D%0A%0A
     

Allow to control FlakyTestRetryAnalyzer with session property