fit mock invocation recorder for targetClass parameter

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

@@ -3,7 +3,6 @@ package com.alibaba.testable.agent.handler;
 import com.alibaba.testable.agent.constant.ConstPool;
 import com.alibaba.testable.agent.model.MethodInfo;
 import com.alibaba.testable.agent.model.ModifiedInsnNodes;
-import com.alibaba.testable.agent.tool.ImmutablePair;
 import com.alibaba.testable.agent.util.BytecodeUtil;
 import com.alibaba.testable.agent.util.ClassUtil;
 import com.alibaba.testable.core.util.LogUtil;
@@ -65,7 +64,7 @@ public class SourceClassHandler extends BaseClassHandler {
         do {
             if (invokeOps.contains(instructions[i].getOpcode())) {
                 MethodInsnNode node = (MethodInsnNode)instructions[i];
-                ImmutablePair<String, String> mockMethod = getMemberInjectMethodName(memberInjectMethods, node);
+                MethodInfo mockMethod = getMemberInjectMethodName(memberInjectMethods, node);
                 if (mockMethod != null) {
                     // it's a member or static method and an inject method for it exist
                     int rangeStart = getMemberMethodStart(instructions, i);
@@ -103,14 +102,13 @@ public class SourceClassHandler extends BaseClassHandler {
      * find the mock method fit for specified method node
      * @param memberInjectMethods mock methods available
      * @param node method node to match for
-     * @return pair of <mock-method-name, mock-method-desc>
+     * @return mock method info
      */
-    private ImmutablePair<String, String> getMemberInjectMethodName(Set<MethodInfo> memberInjectMethods,
+    private MethodInfo getMemberInjectMethodName(Set<MethodInfo> memberInjectMethods, MethodInsnNode node) {
-                                                                    MethodInsnNode node) {
         for (MethodInfo m : memberInjectMethods) {
             String nodeOwner = ClassUtil.fitCompanionClassName(node.owner);
             if (m.getClazz().equals(nodeOwner) && m.getName().equals(node.name) && m.getDesc().equals(node.desc)) {
-                return ImmutablePair.of(m.getMockName(), m.getMockDesc());
+                return m;
             }
         }
         return null;
@@ -219,23 +217,29 @@ public class SourceClassHandler extends BaseClassHandler {
             ClassUtil.toByteCodeClassName(classType);
     }
 
-    private ModifiedInsnNodes replaceMemberCallOps(ClassNode cn, MethodNode mn, ImmutablePair<String, String> mockMethod,
+    private ModifiedInsnNodes replaceMemberCallOps(ClassNode cn, MethodNode mn, MethodInfo mockMethod,
                                                    AbstractInsnNode[] instructions, String ownerClass,
                                                    int opcode, int start, int end) {
-        LogUtil.diagnose("    Line %d, mock method %s used", getLineNum(instructions, start), mockMethod.left);
+        LogUtil.diagnose("    Line %d, mock method %s used", getLineNum(instructions, start),
-        MethodInsnNode method = (MethodInsnNode)instructions[end];
+            mockMethod.getMockName());
+        boolean shouldAppendTypeParameter = !mockMethod.getDesc().equals(mockMethod.getMockDesc());
         String testClassName = ClassUtil.getTestClassName(cn.name);
         if (Opcodes.INVOKESTATIC == opcode || isCompanionMethod(ownerClass, opcode)) {
-            // append a null value if it was a static invoke or in kotlin companion class
+            if (shouldAppendTypeParameter) {
-            mn.instructions.insertBefore(instructions[start], new InsnNode(ACONST_NULL));
+                // append a null value if it was a static invoke or in kotlin companion class
+                mn.instructions.insertBefore(instructions[start], new InsnNode(ACONST_NULL));
+            }
             if (ClassUtil.isCompanionClassName(ownerClass)) {
                 // for kotlin companion class, remove the byte code of reference to "companion" static field
                 mn.instructions.remove(instructions[end - 1]);
             }
+        } else if (!shouldAppendTypeParameter) {
+            // remove extra target ops code
+            mn.instructions.remove(instructions[start]);
         }
         // method with @MockMethod will be modified as public static access, so INVOKESTATIC is used
         mn.instructions.insertBefore(instructions[end], new MethodInsnNode(INVOKESTATIC, testClassName,
-            mockMethod.left, mockMethod.right, false));
+            mockMethod.getMockName(), mockMethod.getMockDesc(), false));
         mn.instructions.remove(instructions[end]);
         return new ModifiedInsnNodes(mn.instructions.toArray(), 1);
     }

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

@@ -4,7 +4,9 @@ import com.alibaba.testable.agent.constant.ConstPool;
 import com.alibaba.testable.agent.tool.ImmutablePair;
 import com.alibaba.testable.agent.util.AnnotationUtil;
 import com.alibaba.testable.agent.util.ClassUtil;
+import com.alibaba.testable.core.model.NullType;
 import com.alibaba.testable.core.util.LogUtil;
+import org.objectweb.asm.Type;
 import org.objectweb.asm.tree.*;
 
 import java.util.List;
@@ -30,7 +32,7 @@ public class TestClassHandler extends BaseClassHandler {
     private static final String METHOD_RECORD_MOCK_INVOKE = "recordMockInvoke";
     private static final String SIGNATURE_CURRENT_TEST_CASE_NAME = "(Ljava/lang/String;)Ljava/lang/String;";
     private static final String SIGNATURE_CURRENT_SOURCE_METHOD_NAME = "()Ljava/lang/String;";
-    private static final String SIGNATURE_INVOKE_RECORDER_METHOD = "([Ljava/lang/Object;Z)V";
+    private static final String SIGNATURE_INVOKE_RECORDER_METHOD = "([Ljava/lang/Object;ZZ)V";
     private static final String SIGNATURE_PARAMETERS = "Ljava/util/Map;";
 
     /**
@@ -193,6 +195,11 @@ public class TestClassHandler extends BaseClassHandler {
         } else {
             il.add(new InsnNode(ICONST_0));
         }
+        if (isTargetClassInParameter(mn)) {
+            il.add(new InsnNode(ICONST_1));
+        } else {
+            il.add(new InsnNode(ICONST_0));
+        }
         il.add(new MethodInsnNode(INVOKESTATIC, CLASS_INVOKE_RECORD_UTIL, METHOD_RECORD_MOCK_INVOKE,
             SIGNATURE_INVOKE_RECORDER_METHOD, false));
         mn.instructions.insertBefore(mn.instructions.get(0), il);
@@ -200,18 +207,32 @@ public class TestClassHandler extends BaseClassHandler {
 
     private boolean isMockForConstructor(MethodNode mn) {
         for (AnnotationNode an : mn.visibleAnnotations) {
-            if (toDotSeparateFullClassName(an.desc).equals(ConstPool.MOCK_CONSTRUCTOR)) {
+            String annotationName = toDotSeparateFullClassName(an.desc);
-                return true;
+            if (ConstPool.MOCK_CONSTRUCTOR.equals(annotationName)) {
-            }
-            String method = AnnotationUtil.getAnnotationParameter
-                (an, ConstPool.FIELD_TARGET_METHOD, null, String.class);
-            if (ConstPool.CONSTRUCTOR.equals(method)) {
                 return true;
+            } else if (ConstPool.MOCK_METHOD.equals(annotationName)) {
+                String method = AnnotationUtil.getAnnotationParameter
+                    (an, ConstPool.FIELD_TARGET_METHOD, null, String.class);
+                if (ConstPool.CONSTRUCTOR.equals(method)) {
+                    return true;
+                }
             }
         }
         return false;
     }
 
+    private boolean isTargetClassInParameter(MethodNode mn) {
+        for (AnnotationNode an : mn.visibleAnnotations) {
+            if (ConstPool.MOCK_METHOD.equals(toDotSeparateFullClassName(an.desc))) {
+                Type type = AnnotationUtil.getAnnotationParameter(an, ConstPool.FIELD_TARGET_CLASS, null, Type.class);
+                if (type != null && !type.getClassName().equals(NullType.class.getName())) {
+                    return false;
+                }
+            }
+        }
+        return true;
+    }
+
     private static ImmutablePair<Integer, Integer> getLoadParameterByteCode(Byte type) {
         switch (type) {
             case ClassUtil.TYPE_BYTE:

testable-agent/src/test/java/com/alibaba/testable/agent/transformer/TestableClassTransformerTest.java

@@ -17,7 +17,7 @@ class TestableClassTransformerTest {
         assertNull(parameters);
         parameters = PrivateAccessor.invoke(testableClassTransformer, "extractFirstParameter", "(Lcom.alibaba.demo.Class;ILjava.lang.String;Z)");
         assertNotNull(parameters);
-        assertEquals("Lcom.alibaba.demo.Class;", parameters.left);
+        assertEquals("com.alibaba.demo.Class", parameters.left);
         assertEquals("(ILjava.lang.String;Z)", parameters.right);
     }
 }

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

@@ -26,7 +26,7 @@ public class InvokeRecordUtil {
      * @param args invocation parameters
      * @param isConstructor whether mocked method is constructor
      */
-    public static void recordMockInvoke(Object[] args, boolean isConstructor) {
+    public static void recordMockInvoke(Object[] args, boolean isConstructor, boolean isTargetClassInParameter) {
         StackTraceElement mockMethodTraceElement = Thread.currentThread().getStackTrace()[INDEX_OF_TEST_CLASS];
         String mockMethodName = mockMethodTraceElement.getMethodName();
         String testClass = mockMethodTraceElement.getClassName();
@@ -37,7 +37,7 @@ public class InvokeRecordUtil {
             records.add(args);
             LogUtil.verbose("  Mock constructor invoked \"%s\"", identify);
         } else {
-            records.add(slice(args, 1));
+            records.add(isTargetClassInParameter ? slice(args, 1) : args);
             LogUtil.verbose("  Mock method invoked \"%s\"", identify);
         }
         INVOKE_RECORDS.put(identify, records);