support enable diagnose log by test class

testable-agent/src/main/java/com/alibaba/testable/agent/PreMain.java

@@ -25,7 +25,7 @@ public class PreMain {
         }
         for (String a : args.split(AND)) {
             if (a.equals(DEBUG)) {
-                LogUtil.enableDebugLog();
+                LogUtil.globalDebugEnable = true;
             }
         }
     }

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

@@ -14,5 +14,6 @@ public class ConstPool {
 
     public static final String FIELD_TARGET_METHOD = "targetMethod";
 
+    public static final String MOCK_WITH = "com.alibaba.testable.core.annotation.MockWith";
     public static final String TESTABLE_MOCK = "com.alibaba.testable.core.annotation.TestableMock";
 }

testable-agent/src/main/java/com/alibaba/testable/agent/handler/SourceClassHandler.java

@@ -52,8 +52,7 @@ public class SourceClassHandler extends BaseClassHandler {
 
     private void transformMethod(ClassNode cn, MethodNode mn, Set<MethodInfo> memberInjectMethods,
                                  Set<MethodInfo> newOperatorInjectMethods) {
-        LogUtil.debug("  Received %d member mock methods, %d constructor mock methods",
+        LogUtil.debug("  Handling method %s", mn.name);
-            memberInjectMethods.size(), newOperatorInjectMethods.size());
         AbstractInsnNode[] instructions = mn.instructions.toArray();
         List<MethodInfo> memberInjectMethodList = new ArrayList<MethodInfo>(memberInjectMethods);
         int i = 0;
@@ -156,7 +155,8 @@ public class SourceClassHandler extends BaseClassHandler {
 
     private AbstractInsnNode[] replaceNewOps(ClassNode cn, MethodNode mn, String newOperatorInjectMethodName,
                                              AbstractInsnNode[] instructions, int start, int end) {
-        LogUtil.debug("  Using %s mock new operation in %s", newOperatorInjectMethodName, mn.name);
+        LogUtil.debug("    Line %d, mock method %s used", getLineNum(instructions, start),
+            newOperatorInjectMethodName);
         String classType = ((TypeInsnNode)instructions[start]).desc;
         String constructorDesc = ((MethodInsnNode)instructions[end]).desc;
         String testClassName = ClassUtil.getTestClassName(cn.name);
@@ -170,6 +170,15 @@ public class SourceClassHandler extends BaseClassHandler {
         return mn.instructions.toArray();
     }
 
+    private int getLineNum(AbstractInsnNode[] instructions, int start) {
+        for (int i = start - 1; i >= 0; i--) {
+            if (instructions[i] instanceof LineNumberNode) {
+                return ((LineNumberNode)instructions[i]).line;
+            }
+        }
+        return 0;
+    }
+
     private String getConstructorInjectDesc(String constructorDesc, String classType) {
         return constructorDesc.substring(0, constructorDesc.length() - 1) +
             ClassUtil.toByteCodeClassName(classType);
@@ -178,7 +187,7 @@ public class SourceClassHandler extends BaseClassHandler {
     private AbstractInsnNode[] replaceMemberCallOps(ClassNode cn, MethodNode mn, String substitutionMethod,
                                                     AbstractInsnNode[] instructions, String ownerClass,
                                                     int opcode, int start, int end) {
-        LogUtil.debug("  Using %s mock method in %s", substitutionMethod, mn.name);
+        LogUtil.debug("    Line %d, mock method %s used", getLineNum(instructions, start), substitutionMethod);
         mn.maxStack++;
         MethodInsnNode method = (MethodInsnNode)instructions[end];
         String testClassName = ClassUtil.getTestClassName(cn.name);

testable-agent/src/main/java/com/alibaba/testable/agent/model/CachedMockParameter.java

@@ -0,0 +1,39 @@
+package com.alibaba.testable.agent.model;
+
+import org.objectweb.asm.tree.AnnotationNode;
+
+/**
+ * Record parameter fetch from @MockWith annotation
+ *
+ * @author flin
+ */
+public class CachedMockParameter {
+
+    private final boolean classExist;
+    private final AnnotationNode mockWith;
+
+    private CachedMockParameter(boolean classExist, AnnotationNode mockWith) {
+        this.classExist = classExist;
+        this.mockWith = mockWith;
+    }
+
+    public static CachedMockParameter notExist() {
+        return new CachedMockParameter(false, null);
+    }
+
+    public static CachedMockParameter exist() {
+        return new CachedMockParameter(true, null);
+    }
+
+    public static CachedMockParameter exist(AnnotationNode mockWith) {
+        return new CachedMockParameter(true, mockWith);
+    }
+
+    public boolean isClassExist() {
+        return classExist;
+    }
+
+    public AnnotationNode getMockWith() {
+        return mockWith;
+    }
+}

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

@@ -3,12 +3,14 @@ package com.alibaba.testable.agent.transformer;
 import com.alibaba.testable.agent.constant.ConstPool;
 import com.alibaba.testable.agent.handler.SourceClassHandler;
 import com.alibaba.testable.agent.handler.TestClassHandler;
+import com.alibaba.testable.agent.model.CachedMockParameter;
 import com.alibaba.testable.agent.tool.ImmutablePair;
 import com.alibaba.testable.agent.model.MethodInfo;
 import com.alibaba.testable.agent.tool.ComparableWeakRef;
 import com.alibaba.testable.agent.util.AnnotationUtil;
 import com.alibaba.testable.agent.util.ClassUtil;
 import com.alibaba.testable.agent.util.LogUtil;
+import com.alibaba.testable.core.model.MockDiagnose;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.tree.AnnotationNode;
 import org.objectweb.asm.tree.ClassNode;
@@ -18,9 +20,7 @@ import java.io.IOException;
 import java.lang.instrument.ClassFileTransformer;
 import java.net.URLClassLoader;
 import java.security.ProtectionDomain;
-import java.util.ArrayList;
+import java.util.*;
-import java.util.List;
-import java.util.Set;
 
 import static com.alibaba.testable.agent.util.ClassUtil.toDotSeparateFullClassName;
 
@@ -29,12 +29,14 @@ import static com.alibaba.testable.agent.util.ClassUtil.toDotSeparateFullClassNa
  */
 public class TestableClassTransformer implements ClassFileTransformer {
 
-    private final Set<ComparableWeakRef<String>> loadedClassNames = ComparableWeakRef.getWeekHashSet();
+    private static final String FIELD_DIAGNOSE = "diagnose";
+    private final Map<ComparableWeakRef<String>, CachedMockParameter> loadedClass =
+        new WeakHashMap<ComparableWeakRef<String>, CachedMockParameter>();
 
     @Override
     public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                             ProtectionDomain protectionDomain, byte[] classFileBuffer) {
-        if (isSystemClass(loader, className) || loadedClassNames.contains(new ComparableWeakRef<String>(className))) {
+        if (isSystemClass(loader, className) || loadedClass.containsKey(new ComparableWeakRef<String>(className))) {
             // Ignore system class and reloaded class
             return null;
         }
@@ -42,15 +44,14 @@ public class TestableClassTransformer implements ClassFileTransformer {
             if (shouldTransformAsSourceClass(className)) {
                 // it's a source class with testable enabled
                 LogUtil.debug("Handling source class %s", className);
-                loadedClassNames.add(new ComparableWeakRef<String>(className));
                 List<MethodInfo> injectMethods = getTestableMockMethods(ClassUtil.getTestClassName(className));
                 return new SourceClassHandler(injectMethods).getBytes(classFileBuffer);
             } else if (shouldTransformAsTestClass(className)) {
                 // it's a test class with testable enabled
                 LogUtil.debug("Handling test class %s", className);
-                loadedClassNames.add(new ComparableWeakRef<String>(className));
                 return new TestClassHandler().getBytes(classFileBuffer);
             }
+            resetMockContext();
         } catch (IOException e) {
             return null;
         }
@@ -58,12 +59,11 @@ public class TestableClassTransformer implements ClassFileTransformer {
     }
 
     private boolean shouldTransformAsSourceClass(String className) {
-        return ClassUtil.anyMethodHasAnnotation(ClassUtil.getTestClassName(className), ConstPool.TESTABLE_MOCK);
+        return hasMockAnnotation(ClassUtil.getTestClassName(className));
     }
 
     private boolean shouldTransformAsTestClass(String className) {
-        return className.endsWith(ConstPool.TEST_POSTFIX) &&
+        return className.endsWith(ConstPool.TEST_POSTFIX) && hasMockAnnotation(className);
-            ClassUtil.anyMethodHasAnnotation(className, ConstPool.TESTABLE_MOCK);
     }
 
     private boolean isSystemClass(ClassLoader loader, String className) {
@@ -107,6 +107,60 @@ public class TestableClassTransformer implements ClassFileTransformer {
         }
     }
 
+    /**
+     * Check whether any method in specified class has specified annotation
+     * @param className class that need to explore
+     * @return found annotation or not
+     */
+    private boolean hasMockAnnotation(String className) {
+        CachedMockParameter cache = loadedClass.get(new ComparableWeakRef<String>(className));
+        if (cache != null) {
+            setupMockContext(cache.getMockWith());
+            return cache.isClassExist();
+        }
+        try {
+            ClassNode cn = new ClassNode();
+            new ClassReader(className).accept(cn, 0);
+            if (cn.visibleAnnotations != null) {
+                for (AnnotationNode an : cn.visibleAnnotations) {
+                    if (toDotSeparateFullClassName(an.desc).equals(ConstPool.MOCK_WITH)) {
+                        setupMockContext(an);
+                        loadedClass.put(new ComparableWeakRef<String>(className), CachedMockParameter.exist(an));
+                        return true;
+                    }
+                }
+            }
+            for (MethodNode mn : cn.methods) {
+                if (mn.visibleAnnotations != null) {
+                    for (AnnotationNode an : mn.visibleAnnotations) {
+                        if (toDotSeparateFullClassName(an.desc).equals(ConstPool.TESTABLE_MOCK)) {
+                            loadedClass.put(new ComparableWeakRef<String>(className), CachedMockParameter.exist());
+                            return true;
+                        }
+                    }
+                }
+            }
+        } catch (Exception e) {
+            // Usually class not found, return without record
+            return false;
+        }
+        loadedClass.put(new ComparableWeakRef<String>(className), CachedMockParameter.notExist());
+        return false;
+    }
+
+    private void setupMockContext(AnnotationNode an) {
+        MockDiagnose mockDebug = AnnotationUtil.getAnnotationParameter(an, FIELD_DIAGNOSE, null, MockDiagnose.class);
+        if (MockDiagnose.ENABLE.equals(mockDebug)) {
+            LogUtil.enableDebugLog();
+        } else if (MockDiagnose.DISABLE.equals(mockDebug)) {
+            LogUtil.disableDebugLog();
+        }
+    }
+
+    private void resetMockContext() {
+        LogUtil.resetDebugLog();
+    }
+
     /**
      * Split desc to "first parameter" and "desc of rest parameters"
      * @param desc method desc

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

@@ -17,9 +17,18 @@ public class AnnotationUtil {
      * @return value of parameter
      */
     public static <T> T getAnnotationParameter(AnnotationNode an, String key, T defaultValue, Class<T> clazz) {
-        if (an.values != null) {
+        if (an != null && an.values != null) {
             for (int i = 0; i < an.values.size(); i += 2) {
                 if (an.values.get(i).equals(key)) {
+                    if (clazz.isEnum()) {
+                        // Enum type are stored as String[] in annotation parameter
+                        String[] values = (String[])an.values.get(i + 1);
+                        if (values == null || values.length != 2) {
+                            return defaultValue;
+                        }
+                        Class<? extends Enum> enumClazz = (Class<? extends Enum>)clazz;
+                        return (T)Enum.valueOf(enumClazz, values[1]);
+                    }
                     return clazz.cast(an.values.get(i + 1));
                 }
             }

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

@@ -43,11 +43,7 @@ public class ClassUtil {
     private static final String EMPTY = "";
     private static final String METHOD_VALUE_OF = "valueOf";
 
-    private final static String JOINER = "::";
-
     private static final Map<Byte, String> TYPE_MAPPING = new HashMap<Byte, String>();
-    private static final Map<ComparableWeakRef<String>, Boolean> loadedClass =
-        new WeakHashMap<ComparableWeakRef<String>, Boolean>();
 
     static {
         TYPE_MAPPING.put(TYPE_BYTE, CLASS_BYTE);
@@ -61,38 +57,6 @@ public class ClassUtil {
         TYPE_MAPPING.put(TYPE_VOID, EMPTY);
     }
 
-    /**
-     * Check whether any method in specified class has specified annotation
-     * @param className class that need to explore
-     * @param annotationName annotation to look for
-     * @return found annotation or not
-     */
-    public static boolean anyMethodHasAnnotation(String className, String annotationName) {
-        String cacheKey = className + JOINER + annotationName;
-        Boolean found = loadedClass.get(new ComparableWeakRef<String>(cacheKey));
-        if (found != null) {
-            return found;
-        }
-        try {
-            ClassNode cn = new ClassNode();
-            new ClassReader(className).accept(cn, 0);
-            for (MethodNode mn : cn.methods) {
-                if (mn.visibleAnnotations != null) {
-                    for (AnnotationNode an : mn.visibleAnnotations) {
-                        if (toDotSeparateFullClassName(an.desc).equals(annotationName)) {
-                            loadedClass.put(new ComparableWeakRef<String>(cacheKey), true);
-                            return true;
-                        }
-                    }
-                }
-            }
-        } catch (Exception e) {
-            // ignore
-        }
-        loadedClass.put(new ComparableWeakRef<String>(cacheKey), false);
-        return false;
-    }
-
     /**
      * fit kotlin companion class name to original name
      * @param name a class name (which could be a companion class)

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

@@ -5,17 +5,31 @@ package com.alibaba.testable.agent.util;
  */
 public class LogUtil {
 
-    private static final int LEVEL_INFO = 0;
+    private static final int LEVEL_ERROR = 0;
+    private static final int LEVEL_WARN = 1;
     private static final int LEVEL_DEBUG = 2;
-    private static int level = LEVEL_INFO;
+    private static final ThreadLocal<Integer> LEVEL = new ThreadLocal<Integer>();
+
+    public static boolean globalDebugEnable = false;
 
     public static void debug(String msg, Object... args) {
-        if (level >= LEVEL_DEBUG) {
+        if (LEVEL.get() >= LEVEL_DEBUG) {
             System.err.println(String.format("[DEBUG] " + msg, args));
         }
     }
 
     public static void enableDebugLog() {
-        level = LEVEL_DEBUG;
+        LEVEL.remove();
+        LEVEL.set(LEVEL_DEBUG);
     }
+
+    public static void disableDebugLog() {
+        LEVEL.remove();
+        LEVEL.set(LEVEL_ERROR);
+    }
+
+    public static void resetDebugLog() {
+        LEVEL.set(globalDebugEnable ? LEVEL_DEBUG : LEVEL_WARN);
+    }
+
 }

testable-agent/src/test/java/com/alibaba/testable/agent/util/ClassUtilTest.java

@@ -2,21 +2,10 @@ package com.alibaba.testable.agent.util;
 
 import org.junit.jupiter.api.Test;
 
-import java.lang.ref.WeakReference;
-import java.util.Map;
-import java.util.WeakHashMap;
-
 import static org.junit.jupiter.api.Assertions.*;
 
 class ClassUtilTest {
 
-    @Test
-    void should_able_to_get_annotation() {
-        assertFalse(ClassUtil.anyMethodHasAnnotation("class.not.exist", ""));
-        assertFalse(ClassUtil.anyMethodHasAnnotation("com.alibaba.testable.agent.util.ClassUtilTest", "annotation.not.exist"));
-        assertTrue(ClassUtil.anyMethodHasAnnotation("com.alibaba.testable.agent.util.ClassUtilTest", "org.junit.jupiter.api.Test"));
-    }
-
     @Test
     void should_able_to_get_parameter_count() {
         assertEquals(0, ClassUtil.getParameterTypes("()V").size());

testable-core/src/main/java/com/alibaba/testable/core/annotation/MockWith.java

@@ -0,0 +1,23 @@
+package com.alibaba.testable.core.annotation;
+
+import com.alibaba.testable.core.model.MockDiagnose;
+
+import java.lang.annotation.*;
+
+/**
+ * Set extra mock parameter to test class
+ *
+ * @author flin
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.TYPE)
+@Documented
+public @interface MockWith {
+
+    /**
+     * switch of mock diagnose information of current test class
+     * @return enable or disable
+     */
+    MockDiagnose diagnose() default MockDiagnose.WARN_ONLY;
+
+}

testable-core/src/main/java/com/alibaba/testable/core/annotation/TestableMock.java

@@ -14,6 +14,7 @@ public @interface TestableMock {
 
     /**
      * mock specified method instead of method with same name
+     * @return target method name
      */
     String targetMethod() default "";
 

testable-core/src/main/java/com/alibaba/testable/core/model/MockDiagnose.java

@@ -0,0 +1,24 @@
+package com.alibaba.testable.core.model;
+
+/**
+ * @author flin
+ */
+
+public enum MockDiagnose {
+
+    /**
+     * Be quiet
+     */
+    DISABLE,
+
+    /**
+     * Only show warning message
+     */
+    WARN_ONLY,
+
+    /**
+     * Print detail diagnose logs
+     */
+    ENABLE
+
+}