always transform all test classes

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

@@ -42,6 +42,7 @@ public class MockClassHandler extends BaseClassWithContextHandler {
 
     @Override
     protected void transform(ClassNode cn) {
+        LogUtil.diagnose("Found mock class %s", cn.name);
         if (!CLASS_OBJECT.equals(cn.superName)) {
             MockAssociationUtil.recordSubMockContainer(ClassUtil.toDotSeparatedName(cn.superName),
                 ClassUtil.toDotSeparatedName(cn.name));

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

@@ -41,6 +41,7 @@ public class SourceClassHandler extends BaseClassHandler {
      */
     @Override
     protected void transform(ClassNode cn) {
+        LogUtil.diagnose("Found source class %s", cn.name);
         Set<MethodInfo> memberInjectMethods = new HashSet<MethodInfo>();
         Set<MethodInfo> newOperatorInjectMethods = new HashSet<MethodInfo>();
         for (MethodInfo im : injectMethods) {

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

@@ -1,12 +1,12 @@
 package com.alibaba.testable.agent.handler;
 
-import com.alibaba.testable.agent.handler.test.*;
+import com.alibaba.testable.agent.handler.test.Framework;
 import com.alibaba.testable.agent.model.TestCaseMethodType;
 import com.alibaba.testable.core.util.LogUtil;
-import org.objectweb.asm.tree.*;
+import org.objectweb.asm.tree.ClassNode;
-
+import org.objectweb.asm.tree.InsnList;
-import java.util.HashSet;
+import org.objectweb.asm.tree.MethodInsnNode;
-import java.util.Set;
+import org.objectweb.asm.tree.MethodNode;
 
 /**
  * @author flin
@@ -20,14 +20,11 @@ public class TestClassHandler extends BaseClassWithContextHandler {
 
     private int testCaseCount = 0;
     private boolean shouldGenerateCleanupMethod = true;
+    private final Framework framework;
 
-    private final Framework[] frameworkClasses = new Framework[] {
+    public TestClassHandler(Framework framework) {
-        new JUnit4Framework(),
+        this.framework = framework;
-        new JUnit5Framework(),
+    }
-        new TestNgFramework(),
-        new TestNgOnClassFramework(),
-        new SpockFramework()
-    };
 
     /**
      * Handle bytecode of test class
@@ -35,11 +32,7 @@ public class TestClassHandler extends BaseClassWithContextHandler {
      */
     @Override
     protected void transform(ClassNode cn) {
-        Framework framework = checkFramework(cn);
+        LogUtil.diagnose("Found test class %s", cn.name);
-        if (framework == null) {
-            LogUtil.warn("Failed to detect test framework for %s", cn.name);
-            return;
-        }
         for (MethodNode mn : cn.methods) {
             handleTestableUtil(mn);
             handleTestCaseMethod(mn, framework);
@@ -52,28 +45,6 @@ public class TestClassHandler extends BaseClassWithContextHandler {
         LogUtil.diagnose("  Found %d test cases", testCaseCount);
     }
 
-    private Framework checkFramework(ClassNode cn) {
-        Set<String> classAnnotationSet = new HashSet<String>();
-        Set<String> methodAnnotationSet = new HashSet<String>();
-        if (cn.visibleAnnotations != null) {
-            for (AnnotationNode an : cn.visibleAnnotations) {
-                classAnnotationSet.add(an.desc);
-            }
-        }
-        for (MethodNode mn : cn.methods) {
-            if (mn.visibleAnnotations != null) {
-                for (AnnotationNode an : mn.visibleAnnotations) {
-                    methodAnnotationSet.add(an.desc);
-                }
-            }
-        }
-        for (Framework i : frameworkClasses) {
-            if (i.fit(classAnnotationSet, methodAnnotationSet)) {
-                return i;
-            }
-        }
-        return null;
-    }
 
     private void handleTestCaseMethod(MethodNode mn, Framework framework) {
         TestCaseMethodType type = framework.checkMethodType(mn);

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

@@ -0,0 +1,47 @@
+package com.alibaba.testable.agent.transformer;
+
+import com.alibaba.testable.agent.handler.test.*;
+import org.objectweb.asm.tree.AnnotationNode;
+import org.objectweb.asm.tree.ClassNode;
+import org.objectweb.asm.tree.MethodNode;
+
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * @author flin
+ */
+public class TestClassChecker {
+
+    private final Framework[] frameworkClasses = new Framework[] {
+        new JUnit4Framework(),
+        new JUnit5Framework(),
+        new TestNgFramework(),
+        new TestNgOnClassFramework(),
+        new SpockFramework()
+    };
+
+    public Framework checkFramework(ClassNode cn) {
+        Set<String> classAnnotationSet = new HashSet<String>();
+        Set<String> methodAnnotationSet = new HashSet<String>();
+        if (cn.visibleAnnotations != null) {
+            for (AnnotationNode an : cn.visibleAnnotations) {
+                classAnnotationSet.add(an.desc);
+            }
+        }
+        for (MethodNode mn : cn.methods) {
+            if (mn.visibleAnnotations != null) {
+                for (AnnotationNode an : mn.visibleAnnotations) {
+                    methodAnnotationSet.add(an.desc);
+                }
+            }
+        }
+        for (Framework i : frameworkClasses) {
+            if (i.fit(classAnnotationSet, methodAnnotationSet)) {
+                return i;
+            }
+        }
+        return null;
+    }
+
+}

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

@@ -5,6 +5,7 @@ import com.alibaba.testable.agent.handler.MockClassHandler;
 import com.alibaba.testable.agent.handler.OmniClassHandler;
 import com.alibaba.testable.agent.handler.SourceClassHandler;
 import com.alibaba.testable.agent.handler.TestClassHandler;
+import com.alibaba.testable.agent.handler.test.Framework;
 import com.alibaba.testable.agent.model.MethodInfo;
 import com.alibaba.testable.agent.util.*;
 import com.alibaba.testable.core.exception.TargetNotExistException;
@@ -42,7 +43,8 @@ public class TestableClassTransformer implements ClassFileTransformer {
     private final String[] BLACKLIST_PREFIXES = new String[] { "sun/", "com/sun/", "javax/crypto/",
         "java/util/logging/", "org/gradle/", "org/robolectric/" };
 
-    public MockClassParser mockClassParser = new MockClassParser();
+    private final MockClassParser mockClassParser = new MockClassParser();
+    private final TestClassChecker testClassChecker = new TestClassChecker();
 
     @Override
     public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
@@ -68,23 +70,26 @@ public class TestableClassTransformer implements ClassFileTransformer {
         try {
             if (mockClassParser.isMockClass(cn)) {
                 // it's a mock class
-                LogUtil.diagnose("Found mock class %s", className);
                 bytes = new MockClassHandler(className).getBytes(bytes);
                 BytecodeUtil.dumpByte(className, GlobalConfig.getDumpPath(), bytes);
-            } else if (foundMockForTestClass(className) != null) {
+                return bytes;
-                // it's a test class with testable enabled
+            }
-                LogUtil.diagnose("Found test class %s", className);
-                bytes = new TestClassHandler().getBytes(bytes);
-                BytecodeUtil.dumpByte(className, GlobalConfig.getDumpPath(), bytes);
-            } else {
             String mockClass = foundMockForSourceClass(className);
             if (mockClass != null) {
                 // it's a source class with testable enabled
                 List<MethodInfo> injectMethods = mockClassParser.getTestableMockMethods(mockClass);
-                    LogUtil.diagnose("Found source class %s", className);
                 bytes = new SourceClassHandler(injectMethods, mockClass).getBytes(bytes);
                 BytecodeUtil.dumpByte(className, GlobalConfig.getDumpPath(), bytes);
+                return bytes;
             }
+            Framework framework = testClassChecker.checkFramework(cn);
+            if (framework != null) {
+                // it's a test class
+                bytes = new TestClassHandler(framework).getBytes(bytes);
+                BytecodeUtil.dumpByte(className, GlobalConfig.getDumpPath(), bytes);
+                return bytes;
+            } else if (cn.name.endsWith(TEST_POSTFIX)) {
+                LogUtil.verbose("Failed to detect test framework for %s", cn.name);
             }
         } catch (TargetNotExistException e) {
             LogUtil.error("Invalid mock method %s::%s - %s", className, e.getMethodName(), e.getMessage());