1package test.tmp.verify;
2
3import org.testng.IMethodInstance;
4import org.testng.IMethodInterceptor;
5import org.testng.ITestContext;
6import org.testng.ITestNGMethod;
7import org.testng.TestNGUtils;
8import org.testng.collections.Maps;
9
10import java.lang.annotation.Annotation;
11import java.lang.reflect.Method;
12import java.util.ArrayList;
13import java.util.List;
14import java.util.Map;
15
16public 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
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   */
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