Update retry extension to handle parameterized tests

Signed-off-by: Hal Spang <[email protected]>

Author: halspang

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