Avoid generating irrelevant imports in test classes

Now methods can request generating additional imports.

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

@@ -21,7 +21,7 @@ import org.jetbrains.annotations.Nullable;
 
 import java.util.Collection;
 
-public class DelegatingTestClassModel implements TestClassModel {
+public class DelegatingTestClassModel extends TestClassModel {
     private final TestClassModel delegate;
 
     public DelegatingTestClassModel(TestClassModel delegate) {
@@ -70,7 +70,7 @@ public class DelegatingTestClassModel implements TestClassModel {
 
     @NotNull
     @Override
-    public Collection<Class<?>> getImports() {
+    public Collection<Class<?>> getOwnImports() {
-        return delegate.getImports();
+        return delegate.getOwnImports();
     }
 }

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

@@ -22,7 +22,7 @@ class RunTestMethodModel(
     }
 
     override fun generateBody(p: Printer) {
-        if (targetBackend == TargetBackend.ANY && additionalRunnerArguments.isEmpty() && testRunnerMethodName == METHOD_NAME) {
+        if (!isWithTargetBackend()) {
             p.println("KotlinTestUtils.$testRunnerMethodName(this::$testMethodName, this, testDataFilePath);")
         } else {
             val className = TargetBackend::class.java.simpleName
@@ -33,6 +33,14 @@ class RunTestMethodModel(
         }
     }
 
+    override fun imports(): Collection<Class<*>> {
+        return super.imports() + if (isWithTargetBackend()) setOf(TargetBackend::class.java) else emptySet()
+    }
+
+    private fun isWithTargetBackend(): Boolean {
+        return !(targetBackend == TargetBackend.ANY && additionalRunnerArguments.isEmpty() && testRunnerMethodName == METHOD_NAME)
+    }
+
     companion object {
         const val METHOD_NAME = "runTest"
     }

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

@@ -29,6 +29,10 @@ class RunTestMethodWithPackageReplacementModel(
         p.println("KotlinTestUtils.$testRunnerMethodName(filePath -> $testMethodName(filePath, packageName), $className.$targetBackend, testDataFilePath$additionalArguments);")
     }
 
+    override fun imports(): Collection<Class<*>> {
+        return super.imports() + setOf(TargetBackend::class.java)
+    }
+
     companion object {
         const val METHOD_NAME = "runTestWithPackageReplacement"
     }

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

@@ -19,7 +19,7 @@ import java.util.*;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
-public class SimpleTestClassModel implements TestClassModel {
+public class SimpleTestClassModel extends TestClassModel {
     private static final Comparator<TestEntityModel> BY_NAME = Comparator.comparing(TestEntityModel::getName);
 
     @NotNull
@@ -238,11 +238,11 @@ public class SimpleTestClassModel implements TestClassModel {
 
     @NotNull
     @Override
-    public Collection<Class<?>> getImports() {
+    public Collection<Class<?>> getOwnImports() {
         return annotations.stream().map(AnnotationModel::getAnnotation).collect(Collectors.toSet());
     }
 
-    private class TestAllFilesPresentMethodModel implements TestMethodModel {
+    private class TestAllFilesPresentMethodModel extends TestMethodModel {
         @NotNull
         @Override
         public String getName() {
@@ -282,11 +282,6 @@ public class SimpleTestClassModel implements TestClassModel {
             return null;
         }
 
-        @Override
-        public void generateSignature(@NotNull Printer p) {
-            TestMethodModel.DefaultImpls.generateSignature(this, p);
-        }
-
         @Override
         public boolean shouldBeGenerated() {
             return true;

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

@@ -20,7 +20,7 @@ import java.util.regex.Pattern;
 
 import static org.jetbrains.kotlin.test.InTextDirectivesUtils.isIgnoredTarget;
 
-public class SimpleTestMethodModel implements TestMethodModel {
+public class SimpleTestMethodModel extends TestMethodModel {
 
     @NotNull
     private final File rootDir;
@@ -98,9 +98,4 @@ public class SimpleTestMethodModel implements TestMethodModel {
         boolean ignored = skipIgnored && isIgnoredTarget(targetBackend, file);
         return (ignored ? "ignore" : "test") + StringsKt.capitalize(TestGeneratorUtil.escapeForJavaIdentifier(unescapedName));
     }
-
-    @Override
-    public void generateSignature(@NotNull Printer p) {
-        TestMethodModel.DefaultImpls.generateSignature(this, p);
-    }
 }

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

@@ -33,7 +33,7 @@ import java.util.List;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
-public class SingleClassTestModel implements TestClassModel {
+public class SingleClassTestModel extends TestClassModel {
     @NotNull
     private final File rootFile;
     @NotNull
@@ -148,11 +148,11 @@ public class SingleClassTestModel implements TestClassModel {
 
     @NotNull
     @Override
-    public Collection<Class<?>> getImports() {
+    public Collection<Class<?>> getOwnImports() {
         return annotations.stream().map(AnnotationModel::getAnnotation).collect(Collectors.toSet());
     }
 
-    private class TestAllFilesPresentMethodModel implements TestMethodModel {
+    private class TestAllFilesPresentMethodModel extends TestMethodModel {
         @NotNull
         @Override
         public String getName() {
@@ -183,11 +183,6 @@ public class SingleClassTestModel implements TestClassModel {
             return null;
         }
 
-        @Override
-        public void generateSignature(@NotNull Printer p) {
-            TestMethodModel.DefaultImpls.generateSignature(this, p);
-        }
-
         @Override
         public boolean shouldBeGenerated() {
             return true;

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

@@ -68,13 +68,13 @@ class TestGenerator(
             p.println("import ", RUNNER.canonicalName, ";")
         }
         p.println("import " + KotlinTestUtils::class.java.canonicalName + ";")
-        p.println("import " + TargetBackend::class.java.canonicalName + ";")
+
-        if (suiteClassPackage != baseTestClassPackage) {
+        for (clazz in testClassModels.flatMapTo(mutableSetOf()) { classModel -> classModel.imports }) {
-            p.println("import $baseTestClassPackage.$baseTestClassName;")
+            p.println("import ${clazz.name};")
         }
 
-        for (clazz in testClassModels.flatMap { classModel -> classModel.imports }.toSet()) {
+        if (suiteClassPackage != baseTestClassPackage) {
-            p.println("import ${clazz.name};")
+            p.println("import $baseTestClassPackage.$baseTestClassName;")
         }
 
         p.println("import " + TestMetadata::class.java.canonicalName + ";")
@@ -98,7 +98,7 @@ class TestGenerator(
                     get() = suiteClassName
             }
         } else {
-            model = object : TestClassModel {
+            model = object : TestClassModel() {
                 override val innerTestClasses: Collection<TestClassModel>
                     get() = testClassModels
 
@@ -120,8 +120,11 @@ class TestGenerator(
                 override val annotations: Collection<AnnotationModel>
                     get() = emptyList()
 
-                override val imports: Collection<Class<*>>
+                override val ownImports: Collection<Class<*>>
                     get() = emptyList()
+
+                override val imports: Collection<Class<*>>
+                    get() = super.imports
             }
         }
 

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

@@ -23,22 +23,35 @@ interface TestEntityModel {
     val dataString: String?
 }
 
-interface TestClassModel : TestEntityModel {
+interface ClassModel : TestEntityModel {
-    val imports: Collection<Class<*>>
+    val ownImports: Collection<Class<*>>
     val innerTestClasses: Collection<TestClassModel>
     val methods: Collection<MethodModel>
     val isEmpty: Boolean
     val dataPathRoot: String?
     val annotations: Collection<AnnotationModel>
+    val imports: Collection<Class<*>>
+}
+
+abstract class TestClassModel : ClassModel {
+    override val imports: Collection<Class<*>>
+        get() {
+            return mutableListOf<Class<*>>().also { allImports ->
+                allImports.addAll(ownImports)
+                methods.flatMapTo(allImports) { it.imports() }
+                innerTestClasses.flatMapTo(allImports) { it.imports }
+            }
+        }
 }
 
 interface MethodModel : TestEntityModel {
     fun shouldBeGenerated(): Boolean = true
     fun generateSignature(p: Printer)
     fun generateBody(p: Printer)
+    fun imports(): Collection<Class<*>> = emptyList()
 }
 
-interface TestMethodModel : MethodModel {
+abstract class TestMethodModel : MethodModel {
     override fun generateSignature(p: Printer) {
         p.print("public void $name() throws Exception")
     }