1 package test.tmp.verify;
2 
3 import org.testng.IMethodInstance;
4 import org.testng.IMethodInterceptor;
5 import org.testng.ITestContext;
6 import org.testng.ITestNGMethod;
7 import org.testng.TestNGUtils;
8 import org.testng.collections.Maps;
9 
10 import java.lang.annotation.Annotation;
11 import java.lang.reflect.Method;
12 import java.util.ArrayList;
13 import java.util.List;
14 import java.util.Map;
15 
16 public class VerifyInterceptor implements IMethodInterceptor {
17 
18   /**
19    * @return the list of methods received in parameters with all methods
20    * annotated with @Verify inserted after each of these test methods.
21    *
22    * This happens in two steps:
23    * - Find all the methods annotated with @Verify in the classes that contain test methods
24    * - Insert these verify methods after each method passed in parameter
25    * These @Verify methods are stored in a map keyed by the class in order to avoid looking them
26    * up more than once on the same class.
27    */
28   @Override
intercept(List<IMethodInstance> methods, ITestContext context)29   public List<IMethodInstance> intercept(List<IMethodInstance> methods,
30       ITestContext context) {
31 
32     List<IMethodInstance> result = new ArrayList<>();
33     Map<Class<?>, List<IMethodInstance>> verifyMethods = Maps.newHashMap();
34     for (IMethodInstance mi : methods) {
35       ITestNGMethod tm = mi.getMethod();
36       List<IMethodInstance> verify = verifyMethods.get(tm.getRealClass());
37       if (verify == null) {
38         verify = findVerifyMethods(tm.getRealClass(), tm);
39       }
40       result.add(mi);
41       result.addAll(verify);
42     }
43 
44     return result;
45   }
46 
47   /**
48    * @return all the @Verify methods found on @code{realClass}
49    */
findVerifyMethods(Class realClass, final ITestNGMethod tm)50   private List<IMethodInstance> findVerifyMethods(Class realClass, final ITestNGMethod tm) {
51     List<IMethodInstance> result = new ArrayList<>();
52     for (final Method m : realClass.getDeclaredMethods()) {
53       Annotation a = m.getAnnotation(Verify.class);
54       if (a != null) {
55         final ITestNGMethod vm = TestNGUtils.createITestNGMethod(tm, m);
56         result.add(new IMethodInstance() {
57 
58           @Override
59           public Object[] getInstances() {
60             return tm.getInstances();
61           }
62 
63           @Override
64           public ITestNGMethod getMethod() {
65             return vm;
66           }
67 
68           public int compareTo(IMethodInstance o) {
69             if (getInstances()[0] == o.getInstances()[0]) {
70               return 0;
71             } else {
72               return -1;
73             }
74           }
75 
76           @Override
77           public Object getInstance() {
78             return tm.getInstance();
79           }
80         });
81       }
82     }
83 
84     return result;
85   }
86 }
87