Update retry extension to handle parameterized tests

Signed-off-by: Hal Spang <halspang@microsoft.com>

Author: halspang

client/src/test/java/com/microsoft/durabletask/ErrorHandlingIntegrationTests.java

@@ -198,7 +198,7 @@ void subOrchestrationException(boolean handleException) throws TimeoutException
     @ValueSource(ints = {1, 2, 10})
     public void retrySubOrchestratorFailures(int maxNumberOfAttempts) throws TimeoutException {
         // There is one task for each sub-orchestrator call and one task between each retry
-        int expectedTaskCount = (maxNumberOfAttempts * 2) - 1;
+        int expectedTaskCount = (maxNumberOfAttempts * 2);
         this.retryOnFailuresCoreTest(maxNumberOfAttempts, expectedTaskCount, ctx -> {
                 RetryPolicy retryPolicy = getCommonRetryPolicy(maxNumberOfAttempts);
                 ctx.callSubOrchestrator(

client/src/test/java/com/microsoft/durabletask/TestRetryExtension.java

@@ -1,33 +1,37 @@
 package com.microsoft.durabletask;
 
 import org.junit.jupiter.api.extension.ExtensionContext;
-import org.junit.jupiter.api.extension.TestExecutionExceptionHandler;
-import org.junit.jupiter.api.extension.BeforeEachCallback;
-import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.InvocationInterceptor;
+import org.junit.jupiter.api.extension.ReflectiveInvocationContext;
+import java.lang.reflect.Method;
 
-public class TestRetryExtension implements TestExecutionExceptionHandler, BeforeEachCallback, AfterEachCallback {
+public class TestRetryExtension implements InvocationInterceptor {
     private static final int MAX_RETRIES = 3;
-    private int currentRetries = 0;
 
     @Override
-    public void beforeEach(ExtensionContext context) {
-        currentRetries = 0;
-    }
-
-    @Override
-    public void handleTestExecutionException(ExtensionContext context, Throwable throwable) throws Throwable {
-        if (currentRetries < MAX_RETRIES - 1) {
-            currentRetries++;
-            System.err.println(String.format("Test '%s' failed on attempt %d/%d", context.getDisplayName(), currentRetries + 1, MAX_RETRIES));
-            context.getRequiredTestMethod().invoke(context.getRequiredTestInstance());
-        } else {
-            System.err.println(String.format("Test '%s' failed after %d attempts", context.getDisplayName(), MAX_RETRIES));
-            throw throwable;
+    public void interceptTestMethod(Invocation<Void> invocation, 
+                                   ReflectiveInvocationContext<Method> invocationContext, 
+                                   ExtensionContext extensionContext) throws Throwable {
+        
+        Throwable lastException = null;
+        
+        for (int attempt = 1; attempt <= MAX_RETRIES; attempt++) {
+            try {
+                invocation.proceed();
+                return; // Success, exit the retry loop
+            } catch (Throwable throwable) {
+                lastException = throwable;
+                if (attempt < MAX_RETRIES) {
+                    System.err.println(String.format("Test '%s' failed on attempt %d/%d: %s", 
+                        extensionContext.getDisplayName(), attempt, MAX_RETRIES, throwable.getMessage()));
+                } else {
+                    System.err.println(String.format("Test '%s' failed after %d attempts", 
+                        extensionContext.getDisplayName(), MAX_RETRIES));
+                }
+            }
         }
-    }
-
-    @Override
-    public void afterEach(ExtensionContext context) {
-        currentRetries = 0;
+        
+        // If we get here, all retries failed
+        throw lastException;
     }
 }