use transmittable thread local to record test case context

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

@@ -2,6 +2,7 @@ package com.alibaba.testable.agent;
 
 import com.alibaba.testable.agent.transformer.TestableClassTransformer;
 import com.alibaba.testable.agent.util.GlobalConfig;
+import com.alibaba.ttl.threadpool.agent.TtlAgent;
 
 import java.lang.instrument.Instrumentation;
 
@@ -20,6 +21,8 @@ public class PreMain {
     public static void premain(String agentArgs, Instrumentation inst) {
         parseArgs(agentArgs);
         inst.addTransformer(new TestableClassTransformer());
+        // add TTL Transformer
+        TtlAgent.premain(agentArgs, inst);
     }
 
     private static void parseArgs(String args) {

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

@@ -12,7 +12,7 @@ abstract public class BaseClassWithContextHandler extends BaseClassHandler {
 
     private static final String CLASS_TESTABLE_TOOL = "com/alibaba/testable/core/tool/TestableTool";
     private static final String CLASS_TESTABLE_UTIL = "com/alibaba/testable/core/util/TestableUtil";
-    private static final String CLASS_MOCK_CONTEXT = "com/alibaba/testable/agent/model/MockContext";
+    private static final String CLASS_MOCK_CONTEXT_HOLDER = "com/alibaba/testable/agent/model/MockContextHolder";
     private static final String FIELD_TEST_CASE = "TEST_CASE";
     private static final String FIELD_SOURCE_METHOD = "SOURCE_METHOD";
     private static final String FIELD_MOCK_CONTEXT = "MOCK_CONTEXT";
@@ -57,7 +57,7 @@ abstract public class BaseClassWithContextHandler extends BaseClassHandler {
             il.add(new MethodInsnNode(INVOKESTATIC, CLASS_TESTABLE_UTIL, METHOD_CURRENT_SOURCE_METHOD_NAME,
                 SIGNATURE_CURRENT_SOURCE_METHOD_NAME, false));
         } else if (FIELD_MOCK_CONTEXT.equals(fieldName)) {
-            il.add(new FieldInsnNode(GETSTATIC, CLASS_MOCK_CONTEXT, FIELD_PARAMETERS, SIGNATURE_PARAMETERS));
+            il.add(new FieldInsnNode(GETSTATIC, CLASS_MOCK_CONTEXT_HOLDER, FIELD_PARAMETERS, SIGNATURE_PARAMETERS));
             il.add(new MethodInsnNode(INVOKESTATIC, CLASS_BASE_CLASS_WITH_CONTEXT_HANDLER, METHOD_GET_TEST_CASE_MARK,
                 SIGNATURE_GET_TEST_CASE_MARK, false));
             il.add(new MethodInsnNode(INVOKEINTERFACE, CLASS_MAP, METHOD_MAP_GET,

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

@@ -1,16 +1,26 @@
 package com.alibaba.testable.agent.handler;
 
-import com.alibaba.testable.core.util.InvokeRecordUtil;
+import com.alibaba.testable.agent.util.ClassUtil;
-import com.alibaba.testable.core.util.TestableUtil;
+import org.objectweb.asm.tree.*;
-import org.objectweb.asm.tree.AbstractInsnNode;
+
-import org.objectweb.asm.tree.ClassNode;
+import java.util.Arrays;
-import org.objectweb.asm.tree.MethodNode;
+import java.util.List;
 
 /**
  * @author flin
  */
 public class TestClassHandler extends BaseClassWithContextHandler {
 
+    private static final String CLASS_MOCK_CONTEXT_UTIL = "com/alibaba/testable/core/util/MockContextUtil";
+    private static final String METHOD_INIT = "init";
+    private static final String DESC_METHOD_INIT = "(Ljava/lang/String;Ljava/lang/String;)V";
+    private final List<String> testAnnotations = Arrays.asList(
+        // JUnit 4
+        "org.junit.Test",
+        // JUnit 5
+        "org.junit.jupiter.api.Test"
+    );
+
     public TestClassHandler(String mockClassName) {
         this.mockClassName = mockClassName;
     }
@@ -23,11 +33,32 @@ public class TestClassHandler extends BaseClassWithContextHandler {
     protected void transform(ClassNode cn) {
         for (MethodNode mn : cn.methods) {
             handleInstruction(cn, mn);
+            handleTestCaseMethod(cn, mn);
         }
     }
 
+    private void handleTestCaseMethod(ClassNode cn, MethodNode mn) {
+        if (mn.visibleAnnotations == null) {
+            return;
+        }
+        for (AnnotationNode an : mn.visibleAnnotations) {
+            if (testAnnotations.contains(ClassUtil.toDotSeparateFullClassName(an.desc))) {
+                injectMockContextInit(cn.name, mn);
+            }
+        }
+    }
+
+    private void injectMockContextInit(String testClassName, MethodNode mn) {
+        InsnList il = new InsnList();
+        il.add(new LdcInsnNode(testClassName));
+        il.add(new LdcInsnNode(mn.name));
+        il.add(new MethodInsnNode(INVOKESTATIC, CLASS_MOCK_CONTEXT_UTIL, METHOD_INIT, DESC_METHOD_INIT, false));
+        mn.instructions.insertBefore(mn.instructions.getFirst(), il);
+    }
+
     private void handleInstruction(ClassNode cn, MethodNode mn) {
         AbstractInsnNode[] instructions = mn.instructions.toArray();
+        // Note: instructions.length will change when instructions updated
         for (int i = 0; i < instructions.length; i++) {
             instructions = handleTestableUtil(cn, mn, instructions, i);
         }

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

@@ -1,12 +1,14 @@
 package com.alibaba.testable.agent.model;
 
+import com.alibaba.testable.agent.tool.UnnullableMap;
+
 import java.util.HashMap;
 import java.util.Map;
 
 /**
  * @author flin
  */
-public class MockContext {
+public class MockContextHolder {
 
     /**
      * Mock method name → ( key → value )

testable-agent/src/main/java/com/alibaba/testable/agent/tool/UnnullableMap.java

@@ -1,4 +1,4 @@
-package com.alibaba.testable.agent.model;
+package com.alibaba.testable.agent.tool;
 
 import com.alibaba.testable.core.accessor.PrivateAccessor;
 

testable-core/pom.xml

@@ -15,9 +15,15 @@
 
     <properties>
         <junit.version>5.6.2</junit.version>
+        <transmittable.thread.local.version>2.12.1</transmittable.thread.local.version>
     </properties>
 
     <dependencies>
+        <dependency>
+            <groupId>com.alibaba</groupId>
+            <artifactId>transmittable-thread-local</artifactId>
+            <version>${transmittable.thread.local.version}</version>
+        </dependency>
         <dependency>
             <groupId>org.junit.jupiter</groupId>
             <artifactId>junit-jupiter-api</artifactId>

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

@@ -0,0 +1,19 @@
+package com.alibaba.testable.core.model;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class MockContext {
+
+    public final String testClassName;
+
+    public final String testCaseName;
+
+    public final Map<String, Object> parameters;
+
+    public MockContext(String testClassName, String testCaseName) {
+        this.testClassName = testClassName;
+        this.testCaseName = testCaseName;
+        this.parameters = new HashMap<String, Object>();
+    }
+}

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

@@ -0,0 +1,28 @@
+package com.alibaba.testable.core.util;
+
+import com.alibaba.testable.core.model.MockContext;
+import com.alibaba.ttl.TransmittableThreadLocal;
+
+public class MockContextUtil {
+
+    public static TransmittableThreadLocal<MockContext> context = new TransmittableThreadLocal<MockContext>();
+
+    /**
+     * Should be invoked at the beginning of each test case method
+     */
+    public static void init(String testClassName, String testCaseName) {
+        context.set(new MockContext(testClassName, testCaseName));
+    }
+
+    /**
+     * Should be invoked at the end of each test case execution
+     */
+    public static void clean() {
+        context.remove();
+    }
+
+    public static String testCaseMark() {
+        MockContext context = MockContextUtil.context.get();
+        return context.testClassName + "::" + context.testCaseName;
+    }
+}