ANTLR和StringTemplate实例：自动生成单元测试类

1. ANTLR语法

要想自动生成单元测试，首先第一步就是分析被测试类。这里以Java代码为例，用ANTLR对Java代码进行分析。要想靠自己完全手写出一门语言的ANTLR语法文件的复杂程度难以想象，很贴心的是在ANTLR的GitHub网站上列出了很多常见语言的语法文件，例如Java，Sqlite和MySQL的SQL语法等。

有了.g4语法文件，按照Antlr v4入门教程和实例中的步骤，就能自动生成出解析器的代码，这里就不再详述了。

2. StringTemplate基础

StringTemplate（简称ST）也是ANTLR提供的一个非常好用的工具。它的功能类似于Velocity、FreeMaker等模板引擎，可以根据事先定义好的模板，在运行时根据不同的传值渲染出不同的网页、邮件、代码等。但从名称中的String就能看出，它是比较轻量级的。试用了一下的确如此，支持在Java中硬编码简单的模板。

但试用过程中还是碰到了有不少问题，总感觉它的模板语法有些复杂啊！而且最新的ST4的API与之前版本发生了很大变化，网上找的很多例子都不好用了。具体还是参考官网的教程吧，以及这个CheatSheet表格。当然，还有花了不少时间调试好的本文的代码示例！

3. 单元测试生成器

首先来看Main方法。输入文本就是code变量表示的一个Java类，里面有两个方法。然后使用自动生成出的ANTLR代码，构建起语义分析器和解析器的处理链，并传入UnitTestGenerator监听器对输入文本进行遍历。

public class JavaCodeParseTest {

    public static void main(String[] args) {
        String code =
                "package com.jcache.store;" +
                "public class CacheStore {" +
                    "Object getCache(int a) {" +
                        "if (a == 1)" +
                            "return 1;" +
                        "else " +
                            "return 2;" +
                    "}" +
                    "void setCache(int a) {" +
                        "return;" +
                    "}" +
                "}";

        // 1.Lexical analysis
        JavaLexer lexer = new JavaLexer(new ANTLRInputStream(code));
        CommonTokenStream tokens = new CommonTokenStream(lexer);

        // 2.Syntax analysis
        JavaParser parser = new JavaParser(tokens);
        ParseTree tree = parser.compilationUnit();

        // 3.Application based on Syntax Tree
        ParseTreeWalker walker = new ParseTreeWalker();
        walker.walk(new UnitTestGenerator(), tree);
    }

}

下面就来看一个核心代码UnitTestGenerator的实现。

这里简单说一下代码中的几个关键点：

• ST模板组：在模板文件中定义比较简单，如果要在Java中定义模板的话，一定要参照本例static初始化块的写法。
• 嵌套子模板和Multi-valued：要想根据方法名列表，自动为每个方法名都生成一个方法，就需要按照

/**
 * Simple unit test generator.
 */
public class UnitTestGenerator extends JavaBaseListener {

    /** Constants: template name, placeholder name, generated code name */
    private static final String CLASS_ST_NAME = "classST";
    private static final String METHOD_ST_NAME = "methodST";

    private static final String TEST_PKG_NAME = "testPkgName";
    private static final String TEST_CLASS_NAME = "testClassName";
    private static final String TEST_METHOD_NAME = "testMethodName";

    private static final String CLASS_NAME_SUFFIX = "Test";
    private static final String METHOD_NAME_PREFIX = "test";

    /** Template for Java */
    private static final String METHOD_ST =
            t("@Test") +
            t("public void " + $(TEST_METHOD_NAME) + "() throws Exception {") +
            tt("//body...") +
            t("}") +
            n("");

    private static final String CLASS_ST =
            n("package " + $(TEST_PKG_NAME) + ";") +
            n("") +
            n("import org.junit.*;") +
            n("") +
            n("public class " + $(TEST_CLASS_NAME) + " {") +
            n("") +
            /**
             * Apply nested template ‘methodST‘ to multi-valued attributes ‘testMethodName‘.
             * NOTE: <attribute:template(argument-list)>
             *      Apply template to attribute with optional argument-list.
             *      Example: <name:bold()> applies bold() to name‘s value.
             *      The first argument of the template gets the iterated value.
             */
            n($(TEST_METHOD_NAME + ":" + METHOD_ST_NAME +"();separator=\"\n\"")) +
            n("}");

    /** ST group to nest template */
    private static STGroup group;
    static {
        group = new STGroup(‘$‘, ‘$‘);

        CompiledST classST = group.defineTemplate(CLASS_ST_NAME, CLASS_ST);
        classST.addArg(new FormalArgument(TEST_PKG_NAME));
        classST.addArg(new FormalArgument(TEST_CLASS_NAME));
        classST.addArg(new FormalArgument(TEST_METHOD_NAME));

        CompiledST methodST = group.defineTemplate(METHOD_ST_NAME, METHOD_ST);
        methodST.addArg(new FormalArgument(TEST_METHOD_NAME));
    }

    /** Attributes. NOTE: group.getInstanceOf() return new ST */
    private Map<String, Object> attributeMap = new HashMap<>();

    @Override
    public void enterPackageDeclaration(@NotNull JavaParser.PackageDeclarationContext ctx) {
        attributeMap.put(TEST_PKG_NAME, ctx.qualifiedName().getText());
    }

    @Override
    public void enterClassDeclaration(@NotNull ClassDeclarationContext ctx) {
        String orgClassName = ctx.Identifier().getText();
        String testClassName = orgClassName + CLASS_NAME_SUFFIX;

        attributeMap.put(TEST_CLASS_NAME, testClassName);
    }

    @Override
    public void enterMethodDeclaration(@NotNull MethodDeclarationContext ctx) {
        String orgMethodName = ctx.Identifier().getText();
        String testMethodName = METHOD_NAME_PREFIX + orgMethodName.substring(0, 1).toUpperCase()
                + orgMethodName.substring(1);

        // Multi-valued attribute
        List<String> methodNames = (List<String>) attributeMap.get(TEST_METHOD_NAME);
        if (methodNames == null) {
            methodNames = new ArrayList<>();
            attributeMap.put(TEST_METHOD_NAME, methodNames);
        }
        methodNames.add(testMethodName);
    }

    @Override
    public void enterStatement(@NotNull StatementContext ctx) {
        // If/else/switch, for/while, try-catch
        switch (ctx.getStart().getText()) {
            case "if":
                break;
            case "for":
                break;
            case "try":
                break;
            default:
                break;
        }
    }

    @Override
    public void exitCompilationUnit(@NotNull JavaParser.CompilationUnitContext ctx) {
        ST template = group.getInstanceOf(CLASS_ST_NAME);
        attributeMap.entrySet().forEach(e -> template.add(e.getKey(), e.getValue()));
        System.out.println(template.render());

        // Cleanup
        attributeMap.clear();
    }

    // =======================================
    //          String Utility
    // =======================================

    private static String $(String attrName) {
        return "$" + attrName + "$";
    }

    private static String n(String str) {
        return str + "\n";
    }

    private static String t(String str) {
        return "\t" + n(str);
    }

    private static String tt(String str) {
        return "\t\t" + n(str);
    }

}