1 package org.testng;
2 
3 import org.testng.collections.Lists;
4 import org.testng.collections.Maps;
5 import org.testng.internal.XmlMethodSelector;
6 
7 import java.util.List;
8 import java.util.Map;
9 import java.util.Set;
10 
11 /**
12  * This class maintains a map of <CODE><Class, List<ITestNGMethod>></CODE>.
13  * It is used by TestWorkers to determine if the method they just ran
14  * is the last of its class, in which case it's time to invoke all the
15  * afterClass methods.
16  *
17  * @author <a href='mailto:the[dot]mindstorm[at]gmail[dot]com'>Alex Popescu</a>
18  */
19 public class ClassMethodMap {
20   private Map<Object, List<ITestNGMethod>> m_classMap = Maps.newHashMap();
21   // These two variables are used throughout the workers to keep track
22   // of what beforeClass/afterClass methods have been invoked
23   private Map<ITestClass, Set<Object>> m_beforeClassMethods = Maps.newHashMap();
24   private Map<ITestClass, Set<Object>> m_afterClassMethods = Maps.newHashMap();
25 
ClassMethodMap(List<ITestNGMethod> methods, XmlMethodSelector xmlMethodSelector)26   public ClassMethodMap(List<ITestNGMethod> methods, XmlMethodSelector xmlMethodSelector) {
27     for (ITestNGMethod m : methods) {
28       // Only add to the class map methods that are included in the
29       // method selector. We can pass a null context here since the selector
30       // should already have been initialized
31       if (xmlMethodSelector != null){
32     	  if (! xmlMethodSelector.includeMethod(null, m, true)) continue;
33       }
34 
35       Object instance = m.getInstance();
36       List<ITestNGMethod> l = m_classMap.get(instance);
37       if (l == null) {
38         l = Lists.newArrayList();
39         m_classMap.put(instance, l);
40       }
41       l.add(m);
42     }
43   }
44 
45   /**
46    * Remove the method from this map and returns true if it is the last
47    * of its class.
48    */
removeAndCheckIfLast(ITestNGMethod m, Object instance)49   public synchronized boolean removeAndCheckIfLast(ITestNGMethod m, Object instance) {
50     List<ITestNGMethod> l = m_classMap.get(instance);
51     if (l != null) {
52       l.remove(m);
53       // It's the last method of this class if all the methods remaining in the list belong to a
54       // different class
55       for (ITestNGMethod tm : l) {
56         if (tm.getEnabled() && tm.getTestClass().equals(m.getTestClass())) return false;
57       }
58       return true;
59     } else {
60       throw new AssertionError("l should not be null");
61     }
62   }
63 
getMethodClass(ITestNGMethod m)64   private Class<?> getMethodClass(ITestNGMethod m) {
65     return m.getTestClass().getRealClass();
66   }
67 
getInvokedBeforeClassMethods()68   public Map<ITestClass, Set<Object>> getInvokedBeforeClassMethods() {
69     return m_beforeClassMethods;
70   }
71 
getInvokedAfterClassMethods()72   public Map<ITestClass, Set<Object>> getInvokedAfterClassMethods() {
73     return m_afterClassMethods;
74   }
75 
clear()76   public void clear() {
77     for(Set<Object> instances: m_beforeClassMethods.values()) {
78       instances.clear();
79       instances= null;
80     }
81     for(Set<Object> instances: m_afterClassMethods.values()) {
82       instances.clear();
83       instances= null;
84     }
85     m_beforeClassMethods.clear();
86     m_afterClassMethods.clear();
87   }
88 }
89