fix NoClassDefFoundError while using MockWith annotation

testable-agent/src/main/java/com/alibaba/testable/agent/constant/ConstPool.java

@@ -12,7 +12,6 @@ public class ConstPool {
 
     public static final String TEST_POSTFIX = "Test";
     public static final String MOCK_POSTFIX = "Mock";
-    public static final String INNER_MOCK_CLASS = "$Mock";
 
     public static final String FIELD_TARGET_METHOD = "targetMethod";
     public static final String FIELD_TARGET_CLASS = "targetClass";

testable-agent/src/main/java/com/alibaba/testable/agent/transformer/TestableClassTransformer.java

@@ -12,6 +12,7 @@ import com.alibaba.testable.agent.util.GlobalConfig;
 import com.alibaba.testable.agent.util.StringUtil;
 import com.alibaba.testable.core.model.MockDiagnose;
 import com.alibaba.testable.core.util.LogUtil;
+import com.alibaba.testable.core.util.MockContextUtil;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.Type;
 import org.objectweb.asm.tree.AnnotationNode;
@@ -43,10 +44,6 @@ public class TestableClassTransformer implements ClassFileTransformer {
     private static final String FIELD_DIAGNOSE = "diagnose";
     private static final String COMMA = ",";
 
-    /**
-     * mock class referred by @MockWith annotation from test or source class
-     */
-    private static final List<String> mockClassesReferred = new ArrayList<String>();
     private static final String CLASS_NAME_MOCK = "Mock";
 
     /**
@@ -128,10 +125,6 @@ public class TestableClassTransformer implements ClassFileTransformer {
         if (mockClass != null) {
             return mockClass;
         }
-        mockClass = ClassUtil.getInnerMockClassName(className);
-        if (isMockClass(mockClass)) {
-            return mockClass;
-        }
         mockClass = ClassUtil.getMockClassName(ClassUtil.getSourceClassName(className));
         if (isMockClass(mockClass)) {
             return mockClass;
@@ -140,7 +133,7 @@ public class TestableClassTransformer implements ClassFileTransformer {
     }
 
     private boolean isMockClass(String className) {
-        return mockClassesReferred.contains(className) || hasMockMethod(className);
+        return MockContextUtil.mockToTests.containsKey(className) || hasMockMethod(className);
     }
 
     private boolean isSystemClass(String className) {
@@ -173,18 +166,16 @@ public class TestableClassTransformer implements ClassFileTransformer {
     }
 
     private List<MethodInfo> getTestableMockMethods(String className) {
-        try {
         List<MethodInfo> methodInfos = new ArrayList<MethodInfo>();
-            ClassNode cn = new ClassNode();
+        ClassNode cn = getClassNode(className);
-            new ClassReader(className).accept(cn, 0);
+        if (cn == null) {
+            return new ArrayList<MethodInfo>();
+        }
         for (MethodNode mn : cn.methods) {
             checkMethodAnnotation(cn, methodInfos, mn);
         }
         LogUtil.diagnose("  Found %d mock methods", methodInfos.size());
         return methodInfos;
-        } catch (Exception e) {
-            return new ArrayList<MethodInfo>();
-        }
     }
 
     private void checkMethodAnnotation(ClassNode cn, List<MethodInfo> methodInfos, MethodNode mn) {
@@ -241,28 +232,11 @@ public class TestableClassTransformer implements ClassFileTransformer {
      * @return name of mock class, null for not found
      */
     private String readMockWithAnnotationAsSourceClass(String className) {
-        try {
+        ClassNode cn = getClassNode(className);
-            ClassNode cn = new ClassNode();
+        if (cn == null) {
-            new ClassReader(className).accept(cn, 0);
-            if (cn.visibleAnnotations != null) {
-                for (AnnotationNode an : cn.visibleAnnotations) {
-                    if (toDotSeparateFullClassName(an.desc).equals(ConstPool.MOCK_WITH)) {
-                        setupDiagnose(an);
-                        if (AnnotationUtil.getAnnotationParameter(an, FIELD_IS_SRC, false, boolean.class)) {
-                            Class<?> value = AnnotationUtil.getAnnotationParameter(an, FIELD_VALUE, null, Class.class);
-                            if (value != null && !NullType.class.equals(value)) {
-                                mockClassesReferred.add(value.getName());
-                                return value.getName();
-                            }
-                        }
-                    }
-                }
-            }
-        } catch (Exception e) {
-            // Usually class not found, return without record
             return null;
         }
-        return null;
+        return parseMockWithAnnotation(cn, true);
     }
 
     /**
@@ -271,32 +245,44 @@ public class TestableClassTransformer implements ClassFileTransformer {
      * @return name of mock class, null for not found
      */
     private String readMockWithAnnotationAndInnerClassAsTestClass(String className) {
-        try {
+        ClassNode cn = getClassNode(className);
-            ClassNode cn = new ClassNode();
+        if (cn == null) {
-            new ClassReader(className).accept(cn, 0);
+            return null;
+        }
+        String mockClassName = parseMockWithAnnotation(cn, false);
+        if (mockClassName != null) {
+            MockContextUtil.mockToTests.get(mockClassName).add(className);
+            return ClassUtil.toSlashSeparatedName(mockClassName);
+        }
+        for (InnerClassNode ic : cn.innerClasses) {
+            if ((ic.access & ACC_PUBLIC) != 0 && (ic.access & ACC_STATIC) != 0 &&
+                ic.name.equals(getInnerMockClassName(className))) {
+                return ic.name;
+            }
+        }
+        return null;
+    }
+
+    /**
+     * Get mock class from @MockWith annotation
+     * @param cn class that may have @MockWith annotation
+     * @return mock class name
+     */
+    private String parseMockWithAnnotation(ClassNode cn, boolean isSrc) {
         if (cn.visibleAnnotations != null) {
             for (AnnotationNode an : cn.visibleAnnotations) {
                 if (toDotSeparateFullClassName(an.desc).equals(ConstPool.MOCK_WITH)) {
                     setupDiagnose(an);
-                        if (!AnnotationUtil.getAnnotationParameter(an, FIELD_IS_SRC, false, boolean.class)) {
+                    if (AnnotationUtil.getAnnotationParameter(an, FIELD_IS_SRC, false, boolean.class) == isSrc) {
-                            Class<?> value = AnnotationUtil.getAnnotationParameter(an, FIELD_VALUE, null, Class.class);
+                        Type clazz = AnnotationUtil.getAnnotationParameter(an, FIELD_VALUE, null, Type.class);
-                            if (value != null && !NullType.class.equals(value)) {
+                        if (clazz == null || NullType.class.getName().equals(clazz.getClassName())) {
-                                mockClassesReferred.add(value.getName());
-                                return value.getName();
-                            }
-                        }
-                    }
-                }
-            }
-            for (InnerClassNode ic : cn.innerClasses) {
-                if ((ic.access & ACC_PUBLIC) != 0 && (ic.access & ACC_STATIC) != 0 && ic.name.equals(CLASS_NAME_MOCK)) {
-                    return ic.name;
-                }
-            }
-        } catch (Exception e) {
-            // Usually class not found, return without record
                             return null;
                         }
+                        return clazz.getClassName();
+                    }
+                }
+            }
+        }
         return null;
     }
 
@@ -306,9 +292,10 @@ public class TestableClassTransformer implements ClassFileTransformer {
      * @return found annotation or not
      */
     private boolean hasMockMethod(String className) {
-        try {
+        ClassNode cn = getClassNode(className);
-            ClassNode cn = new ClassNode();
+        if (cn == null) {
-            new ClassReader(className).accept(cn, 0);
+            return false;
+        }
         for (MethodNode mn : cn.methods) {
             if (mn.visibleAnnotations != null) {
                 for (AnnotationNode an : mn.visibleAnnotations) {
@@ -320,11 +307,21 @@ public class TestableClassTransformer implements ClassFileTransformer {
                 }
             }
         }
-        } catch (Exception e) {
-            // Usually class not found, return without record
         return false;
     }
-        return false;
+
+    private String getInnerMockClassName(String className) {
+        return className + DOLLAR + CLASS_NAME_MOCK;
+    }
+
+    private ClassNode getClassNode(String className) {
+        ClassNode cn = new ClassNode();
+        try {
+            new ClassReader(className).accept(cn, 0);
+        } catch (IOException e) {
+            return null;
+        }
+        return cn;
     }
 
     private void setupDiagnose(AnnotationNode an) {

testable-agent/src/main/java/com/alibaba/testable/agent/util/ClassUtil.java

@@ -84,15 +84,6 @@ public class ClassUtil {
             name.substring(ConstPool.KOTLIN_PREFIX_ACCESS.length()) : name;
     }
 
-    /**
-     * get inner mock class name from test class name
-     * @param testClassName test class name
-     * @return mock class name
-     */
-    public static String getInnerMockClassName(String testClassName) {
-        return testClassName + ConstPool.INNER_MOCK_CLASS;
-    }
-
     /**
      * get mock class name from source class name
      * @param sourceClassName source class name

testable-core/src/main/java/com/alibaba/testable/core/util/MockContextUtil.java

@@ -3,12 +3,17 @@ package com.alibaba.testable.core.util;
 import com.alibaba.testable.core.model.MockContext;
 import com.alibaba.ttl.TransmittableThreadLocal;
 
-import java.util.Map;
+import java.util.*;
 
 public class MockContextUtil {
 
     public static InheritableThreadLocal<MockContext> context = new TransmittableThreadLocal<MockContext>();
 
+    /**
+     * mock class referred by @MockWith annotation to list of its test classes
+     */
+    public static Map<String, Set<String>> mockToTests = UnnullableMap.of(new HashSet<String>());
+
     /**
      * [0]Thread → [1]MockContextUtil → [2]TestClass
      */