FIR IDE: simplify test generation dsl

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

@@ -7,19 +7,21 @@ package org.jetbrains.kotlin.generators
 
 fun generateTestGroupSuiteWithJUnit5(
     args: Array<String>,
+    additionalMethodGenerators: List<MethodGenerator<Nothing>> = emptyList(),
     init: TestGroupSuite.() -> Unit
 ) {
-    generateTestGroupSuiteWithJUnit5(InconsistencyChecker.hasDryRunArg(args), init)
+    generateTestGroupSuiteWithJUnit5(InconsistencyChecker.hasDryRunArg(args), additionalMethodGenerators, init)
 }
 
 fun generateTestGroupSuiteWithJUnit5(
     dryRun: Boolean = false,
+    additionalMethodGenerators: List<MethodGenerator<Nothing>> = emptyList(),
     init: TestGroupSuite.() -> Unit
 ) {
     val suite = TestGroupSuite(ReflectionBasedTargetBackendComputer).apply(init)
     for (testGroup in suite.testGroups) {
         for (testClass in testGroup.testClasses) {
-            val (changed, testSourceFilePath) = NewTestGeneratorImpl.generateAndSave(testClass, dryRun)
+            val (changed, testSourceFilePath) = NewTestGeneratorImpl(additionalMethodGenerators).generateAndSave(testClass, dryRun)
             if (changed) {
                 InconsistencyChecker.inconsistencyChecker(dryRun).add(testSourceFilePath)
             }

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

@@ -29,7 +29,10 @@ private val METHOD_GENERATORS = listOf(
     TransformingTestMethodGenerator,
 )
 
-object NewTestGeneratorImpl : TestGenerator(METHOD_GENERATORS) {
+class NewTestGeneratorImpl(
+    additionalMethodGenerators: List<MethodGenerator<Nothing>>
+) : TestGenerator(METHOD_GENERATORS + additionalMethodGenerators) {
+
     private val GENERATED_FILES = HashSet<String>()
 
     private fun Printer.generateMetadata(testDataSource: TestEntityModel) {
@@ -84,7 +87,7 @@ object NewTestGeneratorImpl : TestGenerator(METHOD_GENERATORS) {
         return generatorInstance.generateAndSave(dryRun)
     }
 
-    private class TestGeneratorInstance(
+    private inner class TestGeneratorInstance(
         baseDir: String,
         suiteTestClassFqName: String,
         baseTestClassFqName: String,
@@ -95,7 +98,8 @@ object NewTestGeneratorImpl : TestGenerator(METHOD_GENERATORS) {
         private val baseTestClassName: String = baseTestClassFqName.substringAfterLast('.', baseTestClassFqName)
         private val suiteClassPackage: String = suiteTestClassFqName.substringBeforeLast('.', baseTestClassPackage)
         private val suiteClassName: String = suiteTestClassFqName.substringAfterLast('.', suiteTestClassFqName)
-        private val testSourceFilePath: String = baseDir + "/" + this.suiteClassPackage.replace(".", "/") + "/" + this.suiteClassName + ".java"
+        private val testSourceFilePath: String =
+            baseDir + "/" + this.suiteClassPackage.replace(".", "/") + "/" + this.suiteClassName + ".java"
 
         init {
             if (!GENERATED_FILES.add(testSourceFilePath)) {
@@ -256,9 +260,11 @@ object NewTestGeneratorImpl : TestGenerator(METHOD_GENERATORS) {
         private fun generateTestMethod(p: Printer, methodModel: MethodModel) {
             val generator = methodGenerators.getValue(methodModel.kind)
 
-            p.generateTestAnnotation()
-            p.generateTags(methodModel)
-            p.generateMetadata(methodModel)
+            if (methodModel.isTestMethod()) {
+                p.generateTestAnnotation()
+                p.generateTags(methodModel)
+                p.generateMetadata(methodModel)
+            }
             generator.hackyGenerateSignature(methodModel, p)
             p.printWithNoIndent(" {")
             p.println()

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

@@ -86,6 +86,11 @@ class TestGroup(
             get() = this@TestGroup.testsRoot
 
         val testModels = ArrayList<TestClassModel>()
+        private val methodModels = mutableListOf<MethodModel>()
+
+        fun method(method: MethodModel) {
+            methodModels += method
+        }
 
         fun model(
             relativeRootPath: String,
@@ -117,14 +122,14 @@ class TestGroup(
                     SingleClassTestModel(
                         rootFile, compiledPattern, compiledExcludedPattern, filenameStartsLowerCase, testMethod, className,
                         realTargetBackend, skipIgnored, testRunnerMethodName, additionalRunnerArguments, annotations,
-                        extractTagsFromDirectory(rootFile)
+                        extractTagsFromDirectory(rootFile), methodModels
                     )
                 } else {
                     SimpleTestClassModel(
                         rootFile, recursive, excludeParentDirs,
                         compiledPattern, compiledExcludedPattern, filenameStartsLowerCase, testMethod, className,
                         realTargetBackend, excludeDirs, skipIgnored, testRunnerMethodName, additionalRunnerArguments, deep, annotations,
-                        extractTagsFromDirectory(rootFile)
+                        extractTagsFromDirectory(rootFile), methodModels
                     )
                 }
             )

generators/test-generator/tests/org/jetbrains/kotlin/generators/model/SimpleTestClassModel.kt

@@ -30,7 +30,8 @@ class SimpleTestClassModel(
     private val additionalRunnerArguments: List<String>,
     private val deep: Int?,
     override val annotations: Collection<AnnotationModel>,
-    override val tags: List<String>
+    override val tags: List<String>,
+    private val additionalMethods: Collection<MethodModel>,
 ) : TestClassModel() {
     override val name: String
         get() = testClassName
@@ -63,7 +64,8 @@ class SimpleTestClassModel(
                         additionalRunnerArguments,
                         if (deep != null) deep - 1 else null,
                         annotations,
-                        extractTagsFromDirectory(file)
+                        extractTagsFromDirectory(file),
+                        additionalMethods.filter { it.shouldBeGeneratedForInnerTestClass() },
                     )
                 )
             }
@@ -100,6 +102,7 @@ class SimpleTestClassModel(
         val result = mutableListOf<MethodModel>()
         result.add(RunTestMethodModel(targetBackend, doTestMethodName, testRunnerMethodName, additionalRunnerArguments))
         result.add(TestAllFilesPresentMethodModel())
+        result.addAll(additionalMethods)
         val listFiles = rootFile.listFiles()
         if (listFiles != null && (deep == null || deep == 0)) {
             for (file in listFiles) {

generators/test-generator/tests/org/jetbrains/kotlin/generators/model/SingleClassTestModel.kt

@@ -23,7 +23,8 @@ class SingleClassTestModel(
     private val testRunnerMethodName: String,
     private val additionalRunnerArguments: List<String>,
     override val annotations: List<AnnotationModel>,
-    override val tags: List<String>
+    override val tags: List<String>,
+    private val additionalMethods: Collection<MethodModel>,
 ) : TestClassModel() {
     override val name: String
         get() = testClassName
@@ -32,6 +33,7 @@ class SingleClassTestModel(
         val result: MutableList<MethodModel> = ArrayList()
         result.add(RunTestMethodModel(targetBackend, doTestMethodName, testRunnerMethodName, additionalRunnerArguments))
         result.add(TestAllFilesPresentMethodModel())
+        result.addAll(additionalMethods)
         FileUtil.processFilesRecursively(rootFile) { file: File ->
             if (!file.isDirectory && filenamePattern.matcher(file.name).matches()) {
                 result.addAll(getTestMethodsFromFile(file))

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

@@ -35,6 +35,8 @@ interface MethodModel : TestEntityModel {
     abstract class Kind
 
     val kind: Kind
+    fun isTestMethod(): Boolean = true
+    fun shouldBeGeneratedForInnerTestClass(): Boolean = true
     fun shouldBeGenerated(): Boolean = true
     fun imports(): Collection<Class<*>> = emptyList()
 }