Allow to add additional annotations to generated classes

generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/AnnotationModel.kt

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.generators.tests.generator
+
+import org.jetbrains.kotlin.test.MuteExtraSuffix
+import org.jetbrains.kotlin.utils.Printer
+
+class AnnotationModel(
+    val annotation: Class<out Annotation>,
+    val arguments: List<Any>
+) {
+    fun generate(p: Printer) {
+        val argumentsString = arguments.joinToString(separator = ",") { "\"$it\""}
+        p.print("@${annotation.simpleName}($argumentsString)")
+    }
+}
+
+fun muteExtraSuffix(suffix: String) = AnnotationModel(MuteExtraSuffix::class.java, arguments = listOf(suffix))
+

generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/DelegatingTestClassModel.java

@@ -61,4 +61,16 @@ public class DelegatingTestClassModel implements TestClassModel {
     public String getDataString() {
         return delegate.getDataString();
     }
+
+    @NotNull
+    @Override
+    public Collection<AnnotationModel> getAnnotations() {
+        return delegate.getAnnotations();
+    }
+
+    @NotNull
+    @Override
+    public Collection<Class<?>> getImports() {
+        return delegate.getImports();
+    }
 }

generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/SimpleTestClassModel.java

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.utils.Printer;
 import java.io.File;
 import java.util.*;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 public class SimpleTestClassModel implements TestClassModel {
     private static final Comparator<TestEntityModel> BY_NAME = Comparator.comparing(TestEntityModel::getName);
@@ -43,6 +44,9 @@ public class SimpleTestClassModel implements TestClassModel {
     @Nullable
     private Collection<MethodModel> testMethods;
 
+    @NotNull
+    private final Collection<AnnotationModel> annotations;
+
     private final boolean skipIgnored;
     private final String testRunnerMethodName;
     private final List<String> additionalRunnerArguments;
@@ -60,7 +64,8 @@ public class SimpleTestClassModel implements TestClassModel {
             boolean skipIgnored,
             String testRunnerMethodName,
             List<String> additionalRunnerArguments,
-            Integer deep
+            Integer deep,
+            @NotNull Collection<AnnotationModel> annotations
     ) {
         this.rootFile = rootFile;
         this.recursive = recursive;
@@ -75,6 +80,7 @@ public class SimpleTestClassModel implements TestClassModel {
         this.testRunnerMethodName = testRunnerMethodName;
         this.additionalRunnerArguments = additionalRunnerArguments;
         this.deep = deep;
+        this.annotations = annotations;
     }
 
     @NotNull
@@ -94,8 +100,9 @@ public class SimpleTestClassModel implements TestClassModel {
                         children.add(new SimpleTestClassModel(
                                 file, true, excludeParentDirs, filenamePattern, checkFilenameStartsLowerCase,
                                 doTestMethodName, innerTestClassName, targetBackend, excludesStripOneDirectory(file.getName()),
-                                skipIgnored, testRunnerMethodName, additionalRunnerArguments, deep != null ? deep - 1 : null)
+                                skipIgnored, testRunnerMethodName, additionalRunnerArguments, deep != null ? deep - 1 : null, annotations)
                         );
+
                     }
                 }
             }
@@ -223,6 +230,18 @@ public class SimpleTestClassModel implements TestClassModel {
         return testClassName;
     }
 
+    @NotNull
+    @Override
+    public Collection<AnnotationModel> getAnnotations() {
+        return annotations;
+    }
+
+    @NotNull
+    @Override
+    public Collection<Class<?>> getImports() {
+        return annotations.stream().map(AnnotationModel::getAnnotation).collect(Collectors.toSet());
+    }
+
     private class TestAllFilesPresentMethodModel implements TestMethodModel {
         @NotNull
         @Override

generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/SingleClassTestModel.java

@@ -31,6 +31,7 @@ import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 public class SingleClassTestModel implements TestClassModel {
     @NotNull
@@ -52,6 +53,9 @@ public class SingleClassTestModel implements TestClassModel {
     private final String testRunnerMethodName;
     private final List<String> additionalRunnerArguments;
 
+    @NotNull
+    private final List<AnnotationModel> annotations;
+
     public SingleClassTestModel(
             @NotNull File rootFile,
             @NotNull Pattern filenamePattern,
@@ -61,7 +65,8 @@ public class SingleClassTestModel implements TestClassModel {
             @NotNull TargetBackend targetBackend,
             boolean skipIgnored,
             String testRunnerMethodName,
-            List<String> additionalRunnerArguments
+            List<String> additionalRunnerArguments,
+            @NotNull List<AnnotationModel> annotations
     ) {
         this.rootFile = rootFile;
         this.filenamePattern = filenamePattern;
@@ -72,6 +77,7 @@ public class SingleClassTestModel implements TestClassModel {
         this.skipIgnored = skipIgnored;
         this.testRunnerMethodName = testRunnerMethodName;
         this.additionalRunnerArguments = additionalRunnerArguments;
+        this.annotations = annotations;
     }
 
     @NotNull
@@ -134,6 +140,18 @@ public class SingleClassTestModel implements TestClassModel {
         return testClassName;
     }
 
+    @NotNull
+    @Override
+    public Collection<AnnotationModel> getAnnotations() {
+        return annotations;
+    }
+
+    @NotNull
+    @Override
+    public Collection<Class<?>> getImports() {
+        return annotations.stream().map(AnnotationModel::getAnnotation).collect(Collectors.toSet());
+    }
+
     private class TestAllFilesPresentMethodModel implements TestMethodModel {
         @NotNull
         @Override

generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/TestGenerationDSL.kt

@@ -15,32 +15,35 @@ class TestGroup(
     private val testsRoot: String,
     val testDataRoot: String,
     val testRunnerMethodName: String,
-    val additionalRunnerArguments: List<String> = emptyList()
+    val additionalRunnerArguments: List<String> = emptyList(),
+    val annotations: List<AnnotationModel> = emptyList()
 ) {
     inline fun <reified T : TestCase> testClass(
         suiteTestClassName: String = getDefaultSuiteTestClassName(T::class.java.simpleName),
         useJunit4: Boolean = false,
+        annotations: List<AnnotationModel> = emptyList(),
         noinline init: TestClass.() -> Unit
     ) {
-        testClass(T::class.java.name, suiteTestClassName, useJunit4, init)
+        testClass(T::class.java.name, suiteTestClassName, useJunit4, annotations, init)
     }
 
     fun testClass(
         baseTestClassName: String,
         suiteTestClassName: String = getDefaultSuiteTestClassName(baseTestClassName.substringAfterLast('.')),
         useJunit4: Boolean,
+        annotations: List<AnnotationModel> = emptyList(),
         init: TestClass.() -> Unit
     ) {
         TestGenerator(
             testsRoot,
             suiteTestClassName,
             baseTestClassName,
-            TestClass().apply(init).testModels,
+            TestClass(annotations).apply(init).testModels,
             useJunit4
         ).generateAndSave()
     }
 
-    inner class TestClass {
+    inner class TestClass(val annotations: List<AnnotationModel>) {
         val testModels = ArrayList<TestClassModel>()
 
         fun model(
@@ -66,13 +69,13 @@ class TestGroup(
                     if (excludeDirs.isNotEmpty()) error("excludeDirs is unsupported for SingleClassTestModel yet")
                     SingleClassTestModel(
                         rootFile, compiledPattern, filenameStartsLowerCase, testMethod, className, targetBackend,
-                        skipIgnored, testRunnerMethodName, additionalRunnerArguments
+                        skipIgnored, testRunnerMethodName, additionalRunnerArguments, annotations
                     )
                 } else {
                     SimpleTestClassModel(
                         rootFile, recursive, excludeParentDirs,
                         compiledPattern, filenameStartsLowerCase, testMethod, className,
-                        targetBackend, excludeDirs, skipIgnored, testRunnerMethodName, additionalRunnerArguments, deep
+                        targetBackend, excludeDirs, skipIgnored, testRunnerMethodName, additionalRunnerArguments, deep, annotations
                     )
                 }
             )

generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/TestGenerator.kt

@@ -25,7 +25,6 @@ class TestGenerator(
     testClassModels: Collection<TestClassModel>,
     useJunit4: Boolean
 ) {
-
     private val baseTestClassPackage: String
     private val suiteClassPackage: String
     private val suiteClassName: String
@@ -73,6 +72,11 @@ class TestGenerator(
         if (suiteClassPackage != baseTestClassPackage) {
             p.println("import $baseTestClassPackage.$baseTestClassName;")
         }
+
+        for (clazz in testClassModels.flatMap { classModel -> classModel.imports }.toSet()) {
+            p.println("import ${clazz.name};")
+        }
+
         p.println("import " + TestMetadata::class.java.canonicalName + ";")
         p.println("import " + RunWith::class.java.canonicalName + ";")
         if (useJunit4) {
@@ -112,6 +116,12 @@ class TestGenerator(
 
                 override val dataPathRoot: String?
                     get() = null
+
+                override val annotations: Collection<AnnotationModel>
+                    get() = emptyList()
+
+                override val imports: Collection<Class<*>>
+                    get() = emptyList()
             }
         }
 
@@ -126,6 +136,8 @@ class TestGenerator(
 
         generateMetadata(p, testClassModel)
         generateTestDataPath(p, testClassModel)
+        generateParameterAnnotations(p, testClassModel)
+
         p.println("@RunWith(", if (useJunit4) JUNIT4_RUNNER.simpleName else RUNNER.simpleName, ".class)")
 
         p.println("public " + staticModifier + "class ", testClassModel.name, " extends ", baseTestClassName, " {")
@@ -201,6 +213,13 @@ class TestGenerator(
             }
         }
 
+        private fun generateParameterAnnotations(p: Printer, testClassModel: TestClassModel) {
+            for (annotationModel in testClassModel.annotations) {
+                annotationModel.generate(p);
+                p.println()
+            }
+        }
+
         private fun generateSuppressAllWarnings(p: Printer) {
             p.println("@SuppressWarnings(\"all\")")
         }

generators/test-generator/tests/org/jetbrains/kotlin/generators/tests/generator/TestModel.kt

@@ -24,10 +24,12 @@ interface TestEntityModel {
 }
 
 interface TestClassModel : TestEntityModel {
+    val imports: Collection<Class<*>>
     val innerTestClasses: Collection<TestClassModel>
     val methods: Collection<MethodModel>
     val isEmpty: Boolean
     val dataPathRoot: String?
+    val annotations: Collection<AnnotationModel>
 }
 
 interface MethodModel : TestEntityModel {