From 0dc9955872424f26ea0cbb1f00a314c598407f9c Mon Sep 17 00:00:00 2001
From: Krishnan Mahadevan <krishnan.mahadevan1978@gmail.com>
Date: Fri, 4 Nov 2022 18:10:45 +0530
Subject: [PATCH] Fixing bug with DataProvider retry

---
 .../java/org/testng/internal/Parameters.java  | 14 +++++++++++
 .../test/dataprovider/DataProviderTest.java   | 12 +++++++++
 .../TestClassFailingRetrySample.java          | 25 +++++++++++++++++++
 3 files changed, 51 insertions(+)
 create mode 100644 testng-core/src/test/java/test/dataprovider/issue2819/TestClassFailingRetrySample.java

diff --git a/testng-core/src/main/java/org/testng/internal/Parameters.java b/testng-core/src/main/java/org/testng/internal/Parameters.java
index f65b8314a0..701de725ec 100644
--- a/testng-core/src/main/java/org/testng/internal/Parameters.java
+++ b/testng-core/src/main/java/org/testng/internal/Parameters.java
@@ -790,6 +790,7 @@ public static ParameterHolder handleParameters(
       }
 
       Iterator<Object[]> initParams = null;
+      RuntimeException thrownException;
       do {
 
         for (IDataProviderListener dataProviderListener : holder.getListeners()) {
@@ -808,18 +809,31 @@ public static ParameterHolder handleParameters(
                   fedInstance,
                   annotationFinder);
           shouldRetry = false;
+          thrownException = null;
         } catch (RuntimeException e) {
           for (IDataProviderListener each : holder.getListeners()) {
             each.onDataProviderFailure(testMethod, methodParams.context, e);
           }
           if (shouldRetry) {
             shouldRetry = retry.retry(dataProviderMethod);
+            thrownException = e;
           } else {
             throw e;
           }
         }
       } while (shouldRetry);
 
+      if (thrownException != null) {
+        // The only time when this will be true is when the following happens:
+        // 1. A Retry was involved with the data provider
+        // 2. The retry mechanism immediately returned false and thus causing
+        // a retry to not happen
+        // 3. Since a retry was not recommended the while loop would have exited but
+        // we still should have been failing the test since the data provider invocation
+        // failed.
+        throw thrownException;
+      }
+
       for (IDataProviderListener dataProviderListener : holder.getListeners()) {
         dataProviderListener.afterDataProviderExecution(
             dataProviderMethod, testMethod, methodParams.context);
diff --git a/testng-core/src/test/java/test/dataprovider/DataProviderTest.java b/testng-core/src/test/java/test/dataprovider/DataProviderTest.java
index e08da4d356..12636d46ff 100644
--- a/testng-core/src/test/java/test/dataprovider/DataProviderTest.java
+++ b/testng-core/src/test/java/test/dataprovider/DataProviderTest.java
@@ -42,6 +42,7 @@
 import test.dataprovider.issue2565.SampleTestUsingSupplier;
 import test.dataprovider.issue2819.DataProviderListenerForRetryAwareTests;
 import test.dataprovider.issue2819.SimpleRetry;
+import test.dataprovider.issue2819.TestClassFailingRetrySample;
 import test.dataprovider.issue2819.TestClassSample;
 import test.dataprovider.issue2819.TestClassUsingDataProviderRetrySample;
 import test.dataprovider.issue2819.TestClassWithMultipleRetryImplSample;
@@ -88,6 +89,17 @@ public void testDataProviderRetryInstancesAreUniqueForEachDataDrivenTest() {
     assertThat(listener.getAfterInvocations()).isEqualTo(2);
   }
 
+  @Test(description = "GITHUB-2819")
+  public void testDataProviderRetryAbortsGracefullyWhenNoRetryAtFirstTime() {
+    TestNG testng = create(TestClassFailingRetrySample.class);
+    DataProviderListenerForRetryAwareTests listener = new DataProviderListenerForRetryAwareTests();
+    testng.addListener(listener);
+    testng.run();
+    assertThat(listener.getBeforeInvocations()).isEqualTo(1);
+    assertThat(listener.getFailureInvocations()).isEqualTo(1);
+    assertThat(listener.getAfterInvocations()).isEqualTo(0);
+  }
+
   @Test(description = "GITHUB-2800")
   public void testDataProviderFromAbstractClassWhenCoupledWithFactories() {
     InvokedMethodNameListener listener = run(test.dataprovider.issue2800.TestClassGenerator.class);
diff --git a/testng-core/src/test/java/test/dataprovider/issue2819/TestClassFailingRetrySample.java b/testng-core/src/test/java/test/dataprovider/issue2819/TestClassFailingRetrySample.java
new file mode 100644
index 0000000000..c90cd8507b
--- /dev/null
+++ b/testng-core/src/test/java/test/dataprovider/issue2819/TestClassFailingRetrySample.java
@@ -0,0 +1,25 @@
+package test.dataprovider.issue2819;
+
+import org.testng.IDataProviderMethod;
+import org.testng.IRetryDataProvider;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+public class TestClassFailingRetrySample {
+
+  @Test(dataProvider = "dp")
+  public void testMethod(int ignored) {}
+
+  @DataProvider(name = "dp", retryUsing = DonotRetry.class)
+  public Object[][] getData() {
+    throw new RuntimeException("problem");
+  }
+
+  public static class DonotRetry implements IRetryDataProvider {
+
+    @Override
+    public boolean retry(IDataProviderMethod dataProvider) {
+      return false;
+    }
+  }
+}