Map object to configurations

CHANGES.txt

@@ -1,7 +1,8 @@
 Current (7.10.0)
-Fixed: GITHUB:3084: Document project's PGP artifact signing keys (Krishnan Mahadevan)
+Fixed: GITHUB-3082: IInvokedMethodListener Iinvoked method does not return correct instance during @BeforeMethod, @AfterMethod and @AfterClass (Krishnan Mahadevan)
+Fixed: GITHUB-3084: Document project's PGP artifact signing keys (Krishnan Mahadevan)
 Fixed: GITHUB-3079: Associate a unique id with every test class object instantiated by TestNG (Krishnan Mahadevan)
-Fixed: GITHUB:3040: replace the usages of synchronized with ReentrantLock (Krishnan Mahadevan)
+Fixed: GITHUB-3040: replace the usages of synchronized with ReentrantLock (Krishnan Mahadevan)
 Fixed: GITHUB-3041: TestNG 7.x DataProvider works in opposite to TestNG 6.x when retrying tests. (Krishnan Mahadevan)
 Fixed: GITHUB-3066: How to dynamically adjust the number of TestNG threads after IExecutorFactory is deprecated? (Krishnan Mahadevan)
 New:   GITHUB-2874: Allow users to define ordering for TestNG listeners (Krishnan Mahadevan)

testng-core/src/main/java/org/testng/ClassMethodMap.java

@@ -32,12 +32,7 @@ public ClassMethodMap(List<ITestNGMethod> methods, XmlMethodSelector xmlMethodSe
       }
 
       Object instance = m.getInstance();
-      Collection<ITestNGMethod> l = classMap.get(instance);
-      if (l == null) {
-        l = new ConcurrentLinkedQueue<>();
-        classMap.put(instance, l);
-      }
-      l.add(m);
+      classMap.computeIfAbsent(instance, k -> new ConcurrentLinkedQueue<>()).add(m);
     }
   }
 

testng-core/src/main/java/org/testng/TestClass.java

@@ -4,6 +4,7 @@
 import java.util.Arrays;
 import java.util.IdentityHashMap;
 import java.util.List;
+import java.util.stream.Stream;
 import org.testng.collections.Lists;
 import org.testng.collections.Objects;
 import org.testng.internal.*;
@@ -32,10 +33,22 @@ class TestClass extends NoOpTestClass implements ITestClass, ITestClassConfigInf
   private final IdentityHashMap<Object, List<ITestNGMethod>> beforeClassConfig =
       new IdentityHashMap<>();
 
+  private final IdentityHashMap<Object, List<ITestNGMethod>> afterClassConfig =
+      new IdentityHashMap<>();
+
   @Override
   public List<ITestNGMethod> getAllBeforeClassMethods() {
-    return beforeClassConfig
-        .values()
+    return getAllClassLevelConfigs(beforeClassConfig);
+  }
+
+  @Override
+  public List<ITestNGMethod> getAllAfterClassMethods() {
+    return getAllClassLevelConfigs(afterClassConfig);
+  }
+
+  private static List<ITestNGMethod> getAllClassLevelConfigs(
+      IdentityHashMap<Object, List<ITestNGMethod>> map) {
+    return map.values()
         .parallelStream()
         .reduce(
             (a, b) -> {
@@ -51,6 +64,11 @@ public List<ITestNGMethod> getInstanceBeforeClassMethods(Object instance) {
     return beforeClassConfig.get(instance);
   }
 
+  @Override
+  public List<ITestNGMethod> getInstanceAfterClassMethods(Object instance) {
+    return afterClassConfig.get(instance);
+  }
+
   private static final Logger LOG = Logger.getLogger(TestClass.class);
 
   protected TestClass(
@@ -201,6 +219,7 @@ private void initMethods() {
               false,
               xmlTest,
               eachInstance);
+      afterClassConfig.put(instance, Arrays.asList(m_afterClassMethods));
       m_beforeGroupsMethods =
           ConfigurationMethod.createBeforeConfigurationMethods(
               objectFactory,
@@ -216,21 +235,29 @@ private void initMethods() {
               false,
               eachInstance);
       m_beforeTestMethods =
-          ConfigurationMethod.createTestMethodConfigurationMethods(
-              objectFactory,
-              testMethodFinder.getBeforeTestMethods(m_testClass),
-              annotationFinder,
-              true,
-              xmlTest,
-              eachInstance);
+          Stream.of(
+                  m_beforeTestMethods,
+                  ConfigurationMethod.createTestMethodConfigurationMethods(
+                      objectFactory,
+                      testMethodFinder.getBeforeTestMethods(m_testClass),
+                      annotationFinder,
+                      true,
+                      xmlTest,
+                      eachInstance))
+              .flatMap(Stream::of)
+              .toArray(ITestNGMethod[]::new);
       m_afterTestMethods =
-          ConfigurationMethod.createTestMethodConfigurationMethods(
-              objectFactory,
-              testMethodFinder.getAfterTestMethods(m_testClass),
-              annotationFinder,
-              false,
-              xmlTest,
-              eachInstance);
+          Stream.of(
+                  m_afterTestMethods,
+                  ConfigurationMethod.createTestMethodConfigurationMethods(
+                      objectFactory,
+                      testMethodFinder.getAfterTestMethods(m_testClass),
+                      annotationFinder,
+                      false,
+                      xmlTest,
+                      eachInstance))
+              .flatMap(Stream::of)
+              .toArray(ITestNGMethod[]::new);
     }
   }
 

testng-core/src/main/java/org/testng/TestRunner.java

@@ -486,15 +486,13 @@ private void initMethods() {
     // Walk through all the TestClasses, store their method
     // and initialize them with the correct ITestClass
     //
+
     for (ITestClass tc : m_classMap.values()) {
       fixMethodsWithClass(tc.getTestMethods(), tc, testMethods);
-      fixMethodsWithClass(
-          ((ITestClassConfigInfo) tc).getAllBeforeClassMethods().toArray(new ITestNGMethod[0]),
-          tc,
-          beforeClassMethods);
+      fixMethodsWithClass(classLevelConfigs(tc, true), tc, beforeClassMethods);
       fixMethodsWithClass(tc.getBeforeTestMethods(), tc, null);
       fixMethodsWithClass(tc.getAfterTestMethods(), tc, null);
-      fixMethodsWithClass(tc.getAfterClassMethods(), tc, afterClassMethods);
+      fixMethodsWithClass(classLevelConfigs(tc, false), tc, afterClassMethods);
       fixMethodsWithClass(tc.getBeforeSuiteMethods(), tc, beforeSuiteMethods);
       fixMethodsWithClass(tc.getAfterSuiteMethods(), tc, afterSuiteMethods);
       fixMethodsWithClass(tc.getBeforeTestConfigurationMethods(), tc, beforeXmlTestMethods);
@@ -558,6 +556,13 @@ private void initMethods() {
             comparator);
   }
 
+  private static ITestNGMethod[] classLevelConfigs(ITestClass tc, boolean beforeConfig) {
+    if (beforeConfig) {
+      return ITestClassConfigInfo.allBeforeClassMethods(tc).toArray(ITestNGMethod[]::new);
+    }
+    return ITestClassConfigInfo.allAfterClassMethods(tc).toArray(ITestNGMethod[]::new);
+  }
+
   private ITestNGMethod[] computeAndGetAllTestMethods() {
     List<ITestNGMethod> testMethods = Lists.newArrayList();
     for (ITestClass tc : m_classMap.values()) {

testng-core/src/main/java/org/testng/internal/ITestClassConfigInfo.java

@@ -1,6 +1,7 @@
 package org.testng.internal;
 
 import java.util.List;
+import org.testng.ITestClass;
 import org.testng.ITestNGMethod;
 import org.testng.collections.Lists;
 
@@ -15,6 +16,10 @@ default List<ITestNGMethod> getAllBeforeClassMethods() {
     return Lists.newArrayList();
   }
 
+  default List<ITestNGMethod> getAllAfterClassMethods() {
+    return Lists.newArrayList();
+  }
+
   /**
    * Query the instance before class methods from config methods map.
    *
@@ -24,4 +29,22 @@ default List<ITestNGMethod> getAllBeforeClassMethods() {
   default List<ITestNGMethod> getInstanceBeforeClassMethods(Object instance) {
     return Lists.newArrayList();
   }
+
+  default List<ITestNGMethod> getInstanceAfterClassMethods(Object instance) {
+    return Lists.newArrayList();
+  }
+
+  static List<ITestNGMethod> allBeforeClassMethods(ITestClass tc) {
+    if (tc instanceof ITestClassConfigInfo) {
+      return ((ITestClassConfigInfo) tc).getAllBeforeClassMethods();
+    }
+    return Lists.newArrayList();
+  }
+
+  static List<ITestNGMethod> allAfterClassMethods(ITestClass tc) {
+    if (tc instanceof ITestClassConfigInfo) {
+      return ((ITestClassConfigInfo) tc).getAllAfterClassMethods();
+    }
+    return Lists.newArrayList();
+  }
 }

testng-core/src/main/java/org/testng/internal/invokers/ConfigInvoker.java

@@ -236,7 +236,7 @@ public void invokeConfigurations(ConfigMethodArguments arguments) {
 
     ITestNGMethod[] methods =
         TestNgMethodUtils.filterMethods(
-            arguments.getTestClass(), arguments.getConfigMethods(), SAME_CLASS);
+            null, arguments.getTestClass(), arguments.getConfigMethods(), SAME_CLASS);
     Object[] parameters = new Object[] {};
 
     for (ITestNGMethod tm : methods) {

testng-core/src/main/java/org/testng/internal/invokers/TestInvoker.java

@@ -183,9 +183,9 @@ public List<ITestResult> invokeTestMethods(
 
     ITestClass testClass = testMethod.getTestClass();
     ITestNGMethod[] beforeMethods =
-        TestNgMethodUtils.filterBeforeTestMethods(testClass, CAN_RUN_FROM_CLASS);
+        TestNgMethodUtils.filterBeforeTestMethods(instance, testClass, CAN_RUN_FROM_CLASS);
     ITestNGMethod[] afterMethods =
-        TestNgMethodUtils.filterAfterTestMethods(testClass, CAN_RUN_FROM_CLASS);
+        TestNgMethodUtils.filterAfterTestMethods(instance, testClass, CAN_RUN_FROM_CLASS);
     int invocationCount = onlyOne ? 1 : testMethod.getInvocationCount();
 
     TestMethodArguments arguments =

testng-core/src/main/java/org/testng/internal/invokers/TestMethodWorker.java

@@ -131,7 +131,7 @@ && doesTaskHavePreRequisites()
 
       // Invoke test method
       try {
-        invokeTestMethods(testMethod, testMethodInstance.getInstance());
+        invokeTestMethods(testMethod, testMethod.getInstance());
       } finally {
         try (KeyAwareAutoCloseableLock.AutoReleasable ignored = lock.lockForObject(key)) {
           invokeAfterClassMethods(testMethod.getTestClass(), testMethodInstance);
@@ -229,10 +229,11 @@ private void invokeAfterClassConfigurations(ITestClass testClass, List<Object> i
       ConfigMethodArguments attributes =
           new Builder()
               .forTestClass(testClass)
-              .usingConfigMethodsAs(testClass.getAfterClassMethods())
               .forSuite(m_testContext.getSuite().getXmlSuite())
               .usingParameters(m_parameters)
               .usingInstance(invokeInstance)
+              .usingConfigMethodsAs(
+                  ((ITestClassConfigInfo) testClass).getInstanceAfterClassMethods(invokeInstance))
               .build();
       m_configInvoker.invokeConfigurations(attributes);
     }

testng-core/src/main/java/org/testng/internal/invokers/TestNgMethodUtils.java

@@ -2,7 +2,9 @@
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.function.BiPredicate;
+import java.util.function.Predicate;
 import org.testng.IClass;
 import org.testng.ITestClass;
 import org.testng.ITestNGMethod;
@@ -69,23 +71,31 @@ private static boolean containsConfigurationMethod(
   }
 
   static ITestNGMethod[] filterBeforeTestMethods(
-      ITestClass testClass, BiPredicate<ITestNGMethod, IClass> predicate) {
-    return filterMethods(testClass, testClass.getBeforeTestMethods(), predicate);
+      Object instance, ITestClass testClass, BiPredicate<ITestNGMethod, IClass> predicate) {
+    return filterMethods(instance, testClass, testClass.getBeforeTestMethods(), predicate);
   }
 
   static ITestNGMethod[] filterAfterTestMethods(
-      ITestClass testClass, BiPredicate<ITestNGMethod, IClass> predicate) {
-    return filterMethods(testClass, testClass.getAfterTestMethods(), predicate);
+      Object instance, ITestClass testClass, BiPredicate<ITestNGMethod, IClass> predicate) {
+    return filterMethods(instance, testClass, testClass.getAfterTestMethods(), predicate);
   }
 
   /** @return Only the ITestNGMethods applicable for this testClass */
   static ITestNGMethod[] filterMethods(
-      IClass testClass, ITestNGMethod[] methods, BiPredicate<ITestNGMethod, IClass> predicate) {
+      Object instance,
+      IClass testClass,
+      ITestNGMethod[] methods,
+      BiPredicate<ITestNGMethod, IClass> predicate) {
     List<ITestNGMethod> vResult = Lists.newArrayList();
+    Predicate<ITestNGMethod> sameInstance =
+        tm ->
+            instance == null
+                || instance.equals(
+                    Optional.ofNullable(tm).map(ITestNGMethod::getInstance).orElse(new Object()));
 
     for (ITestNGMethod tm : methods) {
       String msg;
-      if (predicate.test(tm, testClass)
+      if ((predicate.test(tm, testClass) && sameInstance.test(tm))
           && (!TestNgMethodUtils.containsConfigurationMethod(tm, vResult))) {
         msg = Utils.getVerbose() < 10 ? "" : "Keeping method " + tm + " for class " + testClass;
         vResult.add(tm);

testng-core/src/test/java/test/listeners/ListenersTest.java

@@ -11,6 +11,7 @@
 import java.util.List;
 import java.util.Map;
 import org.assertj.core.api.Assertions;
+import org.assertj.core.api.SoftAssertions;
 import org.testng.ITestNGListener;
 import org.testng.TestNG;
 import org.testng.annotations.DataProvider;
@@ -48,6 +49,8 @@
 import test.listeners.issue2916.TestListenerHolder;
 import test.listeners.issue3064.EvidenceListener;
 import test.listeners.issue3064.SampleTestCase;
+import test.listeners.issue3082.ObjectRepository;
+import test.listeners.issue3082.ObjectTrackingMethodListener;
 
 public class ListenersTest extends SimpleBaseTest {
   public static final String[] github2638ExpectedList =
@@ -521,6 +524,24 @@ public void testSkipStatusInBeforeAndAfterConfigurationMethod(
     assertThat(listener.getLogs()).containsExactlyElementsOf(expected);
   }
 
+  @Test(description = "GITHUB-3082")
+  public void ensureListenerWorksWithCorrectTestClassInstance() {
+    TestNG tng = create(test.listeners.issue3082.TestClassSample.class);
+    tng.addListener(new ObjectTrackingMethodListener());
+    tng.setParallel(XmlSuite.ParallelMode.INSTANCES);
+    tng.setGroupByInstances(true);
+    tng.setVerbose(2);
+    tng.run();
+    assertThat(ObjectRepository.errors()).isEmpty();
+    assertThat(ObjectRepository.instancesCount()).isEqualTo(2);
+    List<String> expected =
+        List.of("beforeClass", "beforeMethod", "testMethod", "afterMethod", "afterClass");
+    SoftAssertions softly = new SoftAssertions();
+    ObjectRepository.invocations()
+        .forEach(it -> softly.assertThat(it).containsExactlyElementsOf(expected));
+    softly.assertAll();
+  }
+
   private void setupTest(boolean addExplicitListener) {
     TestNG testng = new TestNG();
     XmlSuite xmlSuite = createXmlSuite("Xml_Suite");

testng-core/src/test/java/test/listeners/issue3082/IUniqueObject.java

@@ -0,0 +1,8 @@
+package test.listeners.issue3082;
+
+import java.util.UUID;
+
+public interface IUniqueObject {
+
+  UUID id();
+}

testng-core/src/test/java/test/listeners/issue3082/ObjectRepository.java

@@ -0,0 +1,39 @@
+package test.listeners.issue3082;
+
+import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
+
+public final class ObjectRepository {
+
+  private ObjectRepository() {}
+
+  private static final Map<UUID, List<String>> mapping = new ConcurrentHashMap<>();
+
+  private static final Set<UUID> errorObjects = ConcurrentHashMap.newKeySet();
+
+  public static void add(UUID uuid, String methodName) {
+    mapping
+        .computeIfAbsent(uuid, k -> Collections.synchronizedList(new ArrayList<>()))
+        .add(methodName);
+  }
+
+  public static void recordErrorFor(UUID uuid) {
+    errorObjects.add(uuid);
+  }
+
+  public static Set<UUID> errors() {
+    return Collections.unmodifiableSet(errorObjects);
+  }
+
+  public static int instancesCount() {
+    return mapping.keySet().size();
+  }
+
+  public static Collection<List<String>> invocations() {
+    return Collections.unmodifiableCollection(mapping.values());
+  }
+
+  public static boolean hasMapping(UUID uuid, String methodName) {
+    return Optional.ofNullable(mapping.get(uuid)).map(it -> it.contains(methodName)).orElse(false);
+  }
+}

testng-core/src/test/java/test/listeners/issue3082/ObjectTrackingMethodListener.java

@@ -0,0 +1,17 @@
+package test.listeners.issue3082;
+
+import org.testng.IInvokedMethod;
+import org.testng.IInvokedMethodListener;
+import org.testng.ITestResult;
+
+public class ObjectTrackingMethodListener implements IInvokedMethodListener {
+
+  @Override
+  public void beforeInvocation(IInvokedMethod method, ITestResult testResult) {
+    String methodName = method.getTestMethod().getMethodName();
+    Object instance = method.getTestMethod().getInstance();
+    if (instance instanceof IUniqueObject) {
+      ObjectRepository.add(((IUniqueObject) instance).id(), methodName);
+    }
+  }
+}