Examples with RetryOnFailureStrategy org.apache.flink.testutils.junit.extensions.retry.strategy.RetryOnFailureStrategy used on opensource projects

Example 1 with RetryOnFailureStrategy

use of org.apache.flink.testutils.junit.extensions.retry.strategy.RetryOnFailureStrategy in project flink by apache.

the class RetryExtension method provideTestTemplateInvocationContexts.

@Override
public Stream<TestTemplateInvocationContext> provideTestTemplateInvocationContexts(ExtensionContext context) {
    RetryOnFailure retryOnFailure = getRetryAnnotation(context, RetryOnFailure.class);
    RetryOnException retryOnException = getRetryAnnotation(context, RetryOnException.class);
    // sanity check that we don't use both annotations
    if (retryOnFailure != null && retryOnException != null) {
        throw new IllegalArgumentException("You cannot combine the RetryOnFailure and RetryOnException annotations.");
    }
    Map<String, RetryStrategy> testLog = (Map<String, RetryStrategy>) context.getStore(RETRY_NAMESPACE).getOrComputeIfAbsent(RETRY_KEY, key -> new HashMap<>());
    int totalTimes;
    if (retryOnException != null) {
        totalTimes = retryOnException.times() + 1;
        testLog.put(getTestMethodKey(context), new RetryOnExceptionStrategy(totalTimes, retryOnException.exception()));
    } else if (retryOnFailure != null) {
        totalTimes = retryOnFailure.times() + 1;
        testLog.put(getTestMethodKey(context), new RetryOnFailureStrategy(totalTimes));
    } else {
        throw new IllegalArgumentException("Unsupported retry strategy.");
    }
    return IntStream.rangeClosed(1, totalTimes).mapToObj(i -> new RetryContext(i, totalTimes));
}