success to run unit test on private method

src/main/java/com/alibaba/testable/generator/CallSuperMethod.java

@@ -1,9 +1,10 @@
 package com.alibaba.testable.generator;
 
+import com.alibaba.testable.generator.model.Statement;
 import com.alibaba.testable.util.ConstPool;
 import com.alibaba.testable.util.StringUtil;
 import com.sun.tools.javac.tree.JCTree;
-
+import java.lang.reflect.Method;
 import javax.lang.model.element.Modifier;
 import java.util.ArrayList;
 import java.util.List;
@@ -14,38 +15,26 @@ import java.util.List;
  * @author flin
  */
 public class CallSuperMethod {
+
     private final String className;
     private final JCTree.JCMethodDecl method;
-    private Object[] params;
-    private String statement;
 
     public CallSuperMethod(String className, JCTree.JCMethodDecl method) {
         this.className = className;
         this.method = method;
     }
 
-    public Object[] getParams() {
+    public Statement[] invoke() {
-        return params;
+        if (method.getModifiers().getFlags().contains(Modifier.PRIVATE)) {
+            return reflectCall();
+        } else {
+            return commonCall();
+        }
     }
 
-    public String getStatement() {
+    private Statement[] commonCall() {
-        return statement;
-    }
-
-    public CallSuperMethod invoke() {
         List<Object> args = new ArrayList<>();
         StringBuilder code = new StringBuilder();
-        if (method.getModifiers().getFlags().contains(Modifier.PRIVATE)) {
-            reflectCall(args, code);
-        } else {
-            commonCall(args, code);
-        }
-        statement = code.toString();
-        params = args.toArray();
-        return this;
-    }
-
-    private void commonCall(List<Object> args, StringBuilder code) {
         List<String> placeholders = new ArrayList<>();
         for (JCTree.JCVariableDecl p : method.params) {
             args.add(p.name.toString());
@@ -56,21 +45,52 @@ public class CallSuperMethod {
             code.append(".").append(method.name);
         }
         code.append("(").append(StringUtil.join(placeholders, ", ")).append(")");
+        return new Statement[] { returnStatement(new Statement(code.toString(), args.toArray())) };
     }
 
-    private void reflectCall(List<Object> args, StringBuilder code) {
+    private Statement[] reflectCall() {
-        if (!method.restype.toString().equals(ConstPool.CONSTRUCTOR_VOID)) {
+        List<Statement> statements = new ArrayList<>();
-            code.append("(").append(method.restype).append(")");
+        statements.add(getMethodStatement());
+        statements.add(setAccessibleStatement());
+        statements.add(returnStatement(invokeStatement()));
+        return statements.toArray(new Statement[0]);
     }
-        code.append(className).append(".class.getMethod(\"").append(method.name).append("\"");
+
+    private Statement returnStatement(Statement statement) {
+        if (method.restype != null && !method.restype.toString().equals(ConstPool.CONSTRUCTOR_VOID)) {
+            statement.setLine("return " + statement.getLine());
+        }
+        return statement;
+    }
+
+    private Statement getMethodStatement() {
+        List<Object> args = new ArrayList<>();
+        StringBuilder code = new StringBuilder();
+        code.append("$T m = ");
+        args.add(Method.class);
+        code.append(className).append(".class.getDeclaredMethod(\"").append(method.name).append("\"");
         for (JCTree.JCVariableDecl p : method.params) {
             code.append(", $T.class");
             args.add(p.sym.type);
         }
-        code.append(").invoke(this");
+        code.append(")");
+        return new Statement(code.toString(), args.toArray());
+    }
+
+    private Statement setAccessibleStatement() {
+        return new Statement("m.setAccessible(true)", new Object[0]);
+    }
+
+    private Statement invokeStatement() {
+        StringBuilder code = new StringBuilder();
+        if (!method.restype.toString().equals(ConstPool.CONSTRUCTOR_VOID)) {
+            code.append("(").append(method.restype).append(")");
+        }
+        code.append("m.invoke(this");
         for (JCTree.JCVariableDecl p : method.params) {
             code.append(", ").append(p.name);
         }
         code.append(")");
+        return new Statement(code.toString(), new Object[0]);
     }
 }

src/main/java/com/alibaba/testable/generator/model/Statement.java

@@ -0,0 +1,28 @@
+package com.alibaba.testable.generator.model;
+
+public class Statement {
+
+    private String line;
+    private Object[] params;
+
+    public String getLine() {
+        return line;
+    }
+
+    public void setLine(String line) {
+        this.line = line;
+    }
+
+    public Object[] getParams() {
+        return params;
+    }
+
+    public void setParams(Object[] params) {
+        this.params = params;
+    }
+
+    public Statement(String line, Object[] params) {
+        this.line = line;
+        this.params = params;
+    }
+}

src/main/java/com/alibaba/testable/processor/TestableProcessor.java

@@ -2,7 +2,9 @@ package com.alibaba.testable.processor;
 
 import com.alibaba.testable.annotation.Testable;
 import com.alibaba.testable.generator.CallSuperMethod;
-import com.alibaba.testable.translator.TestableTreeTranslator;
+import com.alibaba.testable.generator.model.Statement;
+import com.alibaba.testable.translator.TestableClassTranslator;
+import com.alibaba.testable.translator.TestableFieldTranslator;
 import com.alibaba.testable.util.ConstPool;
 import com.squareup.javapoet.*;
 import com.sun.tools.javac.code.Type;
@@ -32,55 +34,71 @@ public class TestableProcessor extends BaseProcessor {
 
     @Override
     public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
-        Set<? extends Element> clazz = roundEnv.getElementsAnnotatedWith(Testable.class);
+        Set<? extends Element> elements = roundEnv.getElementsAnnotatedWith(Testable.class);
-        for (Element classElement : clazz) {
+        for (Element element : elements) {
-            String packageName = elementUtils.getPackageOf(classElement).getQualifiedName().toString();
+            if (element.getKind().isClass()) {
-            String testableTypeName = classElement.getSimpleName().toString().replace(".", "_") + "Testable";
+                processClassElement(element);
-            String fullQualityTypeName =  packageName + "." + testableTypeName;
+            } else if (element.getKind().isField()) {
-            try {
+                processFieldElement(element);
-                JavaFileObject jfo = filter.createSourceFile(fullQualityTypeName);
-                Writer writer = jfo.openWriter();
-                writer.write(createTestableClass(classElement, packageName, testableTypeName));
-                writer.close();
-            } catch (IOException e) {
-                e.printStackTrace();
             }
         }
         return true;
     }
 
-    private String createTestableClass(Element classElement, String packageName, String className) {
+    private void processFieldElement(Element field) {
+        JCTree tree = trees.getTree(field);
+        tree.accept(new TestableFieldTranslator(treeMaker));
+    }
 
-        JCTree tree = trees.getTree(classElement);
+    private void processClassElement(Element clazz) {
-        TestableTreeTranslator translator = new TestableTreeTranslator();
+        String packageName = elementUtils.getPackageOf(clazz).getQualifiedName().toString();
+        String testableTypeName = clazz.getSimpleName().toString().replace(".", "_") + "Testable";
+        String fullQualityTypeName =  packageName + "." + testableTypeName;
+        try {
+            JavaFileObject jfo = filter.createSourceFile(fullQualityTypeName);
+            Writer writer = jfo.openWriter();
+            writer.write(createTestableClass(clazz, packageName, testableTypeName));
+            writer.close();
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
+    }
+
+    private String createTestableClass(Element clazz, String packageName, String className) {
+        JCTree tree = trees.getTree(clazz);
+        TestableClassTranslator translator = new TestableClassTranslator();
         tree.accept(translator);
 
         List<MethodSpec> methodSpecs = new ArrayList<>();
         for (JCTree.JCMethodDecl method : translator.getMethods()) {
-            if (method.getModifiers().getFlags().contains(Modifier.ABSTRACT)) {
+            if (isNoncallableMethod(method)) {
                 continue;
             }
-            if (method.name.toString().equals(ConstPool.CONSTRUCTOR_NAME)) {
+            if (isConstructorMethod(method)) {
-                MethodSpec.Builder builder = MethodSpec.constructorBuilder()
+                buildConstructorMethod(clazz, methodSpecs, method);
-                    .addModifiers(Modifier.PUBLIC);
-                for (JCTree.JCVariableDecl p : method.getParameters()) {
-                    builder.addParameter(getParameterSpec(p));
-                }
-                CallSuperMethod callSuperMethod = new CallSuperMethod(classElement.getSimpleName().toString(), method).invoke();
-                builder.addStatement(callSuperMethod.getStatement(), callSuperMethod.getParams());
-                methodSpecs.add(builder.build());
             } else {
+                buildMemberMethod(clazz, methodSpecs, method);
+            }
+        }
+
+        TypeSpec.Builder builder = TypeSpec.classBuilder(className)
+            .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
+            .superclass(clazz.asType());
+        for (MethodSpec m : methodSpecs) {
+            builder.addMethod(m);
+        }
+        TypeSpec testableClass = builder.build();
+        JavaFile javaFile = JavaFile.builder(packageName, testableClass).build();
+        return javaFile.toString();
+    }
+
+    private void buildMemberMethod(Element classElement, List<MethodSpec> methodSpecs, JCTree.JCMethodDecl method) {
         MethodSpec.Builder builder = MethodSpec.methodBuilder(method.name.toString())
             .addModifiers(toPublicFlags(method.getModifiers()))
             .returns(TypeName.get(((Type.MethodType)method.sym.type).restype));
         for (JCTree.JCVariableDecl p : method.getParameters()) {
             builder.addParameter(getParameterSpec(p));
         }
-                CallSuperMethod callSuperMethod = new CallSuperMethod(classElement.getSimpleName().toString(), method).invoke();
-                String statement = callSuperMethod.getStatement();
-                if (!method.restype.toString().equals(ConstPool.CONSTRUCTOR_VOID)) {
-                    statement = "return " + statement;
-                }
         if (method.getModifiers().getFlags().contains(Modifier.PRIVATE)) {
             builder.addException(Exception.class);
         } else {
@@ -89,20 +107,34 @@ public class TestableProcessor extends BaseProcessor {
                 builder.addException(TypeName.get(exception.type));
             }
         }
-                builder.addStatement(statement, callSuperMethod.getParams());
+        addStatements(builder, classElement, method);
         methodSpecs.add(builder.build());
     }
+
+    private void buildConstructorMethod(Element classElement, List<MethodSpec> methodSpecs,
+                                        JCTree.JCMethodDecl method) {
+        MethodSpec.Builder builder = MethodSpec.constructorBuilder()
+            .addModifiers(Modifier.PUBLIC);
+        for (JCTree.JCVariableDecl p : method.getParameters()) {
+            builder.addParameter(getParameterSpec(p));
+        }
+        addStatements(builder, classElement, method);
+        methodSpecs.add(builder.build());
+    }
+
+    private void addStatements(MethodSpec.Builder builder, Element classElement, JCTree.JCMethodDecl method) {
+        Statement[] statements = new CallSuperMethod(classElement.getSimpleName().toString(), method).invoke();
+        for (Statement s : statements) {
+            builder.addStatement(s.getLine(), s.getParams());
+        }
     }
 
-        TypeSpec.Builder builder = TypeSpec.classBuilder(className)
+    private boolean isConstructorMethod(JCTree.JCMethodDecl method) {
-            .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
+        return method.name.toString().equals(ConstPool.CONSTRUCTOR_NAME);
-            .superclass(classElement.asType());
-        for (MethodSpec m : methodSpecs) {
-            builder.addMethod(m);
     }
-        TypeSpec testableClass = builder.build();
+
-        JavaFile javaFile = JavaFile.builder(packageName, testableClass).build();
+    private boolean isNoncallableMethod(JCTree.JCMethodDecl method) {
-        return javaFile.toString();
+        return method.getModifiers().getFlags().contains(Modifier.ABSTRACT);
     }
 
     private Set<Modifier> toPublicFlags(JCTree.JCModifiers modifiers) {

src/main/java/com/alibaba/testable/translator/TestableClassTranslator.java

@@ -12,7 +12,7 @@ import java.lang.reflect.Modifier;
  *
  * @author flin
  */
-public class TestableTreeTranslator extends TreeTranslator {
+public class TestableClassTranslator extends TreeTranslator {
 
     /**
      * Methods to inject

src/main/java/com/alibaba/testable/translator/TestableFieldTranslator.java

@@ -0,0 +1,36 @@
+package com.alibaba.testable.translator;
+
+import com.sun.tools.javac.tree.JCTree;
+import com.sun.tools.javac.tree.TreeMaker;
+import com.sun.tools.javac.tree.TreeTranslator;
+import com.sun.tools.javac.util.Name;
+
+/**
+ * Travel AST
+ *
+ * @author flin
+ */
+public class TestableFieldTranslator extends TreeTranslator {
+
+    private TreeMaker treeMaker;
+
+    public TestableFieldTranslator(TreeMaker treeMaker) {this.treeMaker = treeMaker;}
+
+    @Override
+    public void visitVarDef(JCTree.JCVariableDecl decl) {
+        super.visitVarDef(decl);
+        decl.vartype = getTestableClassIdent(decl.vartype);
+    }
+
+    @Override
+    public void visitNewClass(JCTree.JCNewClass jcNewClass) {
+        super.visitNewClass(jcNewClass);
+        jcNewClass.clazz = getTestableClassIdent(jcNewClass.clazz);
+    }
+
+    private JCTree.JCIdent getTestableClassIdent(JCTree.JCExpression clazz) {
+        Name className = ((JCTree.JCIdent)clazz).name;
+        return treeMaker.Ident(className.table.fromString(className + "Testable"));
+    }
+
+}