Select Kotlin test src root by default when generating for Kotlin #949 (#1074)

Author: volivan239

utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/generator/CodeGenerationController.kt

@@ -416,8 +416,8 @@ object CodeGenerationController {
         // all test roots for the given test module
         val testRoots = runReadAction {
             testModule
-                .suitableTestSourceRoots(this)
-                .mapNotNull { psiManager.findDirectory(it) }
+                .suitableTestSourceRoots()
+                .mapNotNull { psiManager.findDirectory(it.dir) }
         }
 
         // return an util class from one of the test source roots or null if no util class was found

utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/models/GenerateTestsModel.kt

@@ -19,6 +19,7 @@ import com.intellij.openapi.vfs.newvfs.impl.FakeVirtualFile
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiJavaFile
 import com.intellij.refactoring.util.classMembers.MemberInfo
+import org.jetbrains.kotlin.asJava.classes.KtUltraLightClass
 import org.jetbrains.kotlin.psi.KtFile
 import org.utbot.framework.plugin.api.JavaDocCommentStyle
 import org.utbot.framework.util.ConflictTriggers
@@ -56,13 +57,14 @@ data class GenerateTestsModel(
             ?: error("Could not find module for $newTestSourceRoot")
     }
 
+    var codegenLanguage = if (srcClasses.all { it is KtUltraLightClass }) CodegenLanguage.KOTLIN else CodegenLanguage.JAVA
+
     var testPackageName: String? = null
     lateinit var testFramework: TestFramework
     lateinit var mockStrategy: MockStrategyApi
     lateinit var mockFramework: MockFramework
     lateinit var staticsMocking: StaticsMocking
     lateinit var parametrizedTestSource: ParametrizedTestSource
-    lateinit var codegenLanguage: CodegenLanguage
     lateinit var runtimeExceptionTestsBehaviour: RuntimeExceptionTestsBehaviour
     lateinit var hangingTestsTimeout: HangingTestsTimeout
     lateinit var forceStaticMocking: ForceStaticMocking

utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/GenerateTestsDialogWindow.kt

@@ -95,7 +95,6 @@ import javax.swing.JSpinner
 import javax.swing.text.DefaultFormatter
 import org.jetbrains.concurrency.Promise
 import org.jetbrains.concurrency.thenRun
-import org.jetbrains.kotlin.asJava.classes.KtUltraLightClass
 import org.utbot.common.PathUtil.toPath
 import org.utbot.framework.UtSettings
 import org.utbot.framework.codegen.ForceStaticMocking
@@ -651,8 +650,7 @@ class GenerateTestsDialogWindow(val model: GenerateTestsModel) : DialogWrapper(m
         mockStrategies.isEnabled = areMocksSupported
         staticsMocking.isEnabled = areMocksSupported && mockStrategies.item != MockStrategyApi.NO_MOCKS
 
-        codegenLanguages.item =
-            if (model.srcClasses.all { it is KtUltraLightClass }) CodegenLanguage.KOTLIN else CodegenLanguage.JAVA
+        codegenLanguages.item = model.codegenLanguage
 
         val installedTestFramework = TestFramework.allItems.singleOrNull { it.isInstalled }
         currentFrameworkItem = when (parametrizedTestSources.isSelected) {

utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/components/TestFolderComboWithBrowseButton.kt

@@ -19,6 +19,7 @@ import javax.swing.JList
 import org.jetbrains.kotlin.idea.util.projectStructure.allModules
 import org.utbot.common.PathUtil
 import org.utbot.intellij.plugin.models.GenerateTestsModel
+import org.utbot.intellij.plugin.ui.utils.TestSourceRoot
 import org.utbot.intellij.plugin.ui.utils.addDedicatedTestRoot
 import org.utbot.intellij.plugin.ui.utils.isBuildWithGradle
 import org.utbot.intellij.plugin.ui.utils.suitableTestSourceRoots
@@ -58,9 +59,20 @@ class TestFolderComboWithBrowseButton(private val model: GenerateTestsModel) :
         val suggestedModules =
             if (model.project.isBuildWithGradle) model.project.allModules() else model.potentialTestModules
 
-        val testRoots = suggestedModules.flatMap { it.suitableTestSourceRoots().toList() }.toMutableList()
+        val testRoots = suggestedModules.flatMap {
+            it.suitableTestSourceRoots()
+        }.sortedWith(
+            compareByDescending<TestSourceRoot> {
+                // Heuristics: Dirs with language == codegenLanguage should go first
+                it.expectedLanguage == model.codegenLanguage
+            }.thenBy {
+                // Heuristics: User is more likely to choose the shorter path
+                it.dir.path.length
+            }
+        ).toMutableList()
+
         // this method is blocked for Gradle, where multiple test modules can exist
-        model.testModule.addDedicatedTestRoot(testRoots)
+        model.testModule.addDedicatedTestRoot(testRoots, model.codegenLanguage)
 
         if (testRoots.isNotEmpty()) {
             configureRootsCombo(testRoots)
@@ -94,13 +106,12 @@ class TestFolderComboWithBrowseButton(private val model: GenerateTestsModel) :
             files.singleOrNull()
         }
 
-    private fun configureRootsCombo(testRoots: List<VirtualFile>) {
-        // unfortunately, Gradle creates Kotlin test source root with Java source root type, so type is misleading
+    private fun configureRootsCombo(testRoots: List<TestSourceRoot>) {
         val selectedRoot = testRoots.first()
 
         // do not update model.testModule here, because fake test source root could have been chosen
-        model.testSourceRoot = selectedRoot
-        newItemList(testRoots.toSet())
+        model.testSourceRoot = selectedRoot.dir
+        newItemList(testRoots.map { it.dir }.toSet())
     }
 
     private fun newItemList(comboItems: Set<Any>) {

utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/utils/ModuleUtils.kt

@@ -37,6 +37,11 @@ import org.jetbrains.kotlin.platform.TargetPlatformVersion
 
 private val logger = KotlinLogging.logger {}
 
+data class TestSourceRoot(
+    val dir: VirtualFile,
+    val expectedLanguage: CodegenLanguage
+)
+
 /**
  * @return jdk version of the module
  */
@@ -60,12 +65,6 @@ fun Module.kotlinTargetPlatform(): TargetPlatformVersion {
         ?.singleOrNull() ?: error("Can't determine target platform for module $this")
 }
 
-fun Module.suitableTestSourceRoots(): List<VirtualFile> =
-    suitableTestSourceRoots(CodegenLanguage.JAVA) + suitableTestSourceRoots(CodegenLanguage.KOTLIN)
-
-fun Module.suitableTestSourceFolders(): List<SourceFolder> =
-    suitableTestSourceFolders(CodegenLanguage.JAVA) + suitableTestSourceFolders(CodegenLanguage.KOTLIN)
-
 /**
  * Gets a path to test resources source root.
  *
@@ -121,8 +120,8 @@ private fun findPotentialModulesForTests(project: Project, srcModule: Module): L
 /**
  * Finds all suitable test root virtual files.
  */
-fun Module.suitableTestSourceRoots(codegenLanguage: CodegenLanguage): List<VirtualFile> {
-    val sourceRootsInModule = suitableTestSourceFolders(codegenLanguage).mapNotNull { it.file }
+fun Module.suitableTestSourceRoots(): List<TestSourceRoot> {
+    val sourceRootsInModule = suitableTestSourceFolders().mapNotNull { it.testSourceRoot }
 
     if (sourceRootsInModule.isNotEmpty()) {
         return sourceRootsInModule
@@ -133,22 +132,30 @@ fun Module.suitableTestSourceRoots(codegenLanguage: CodegenLanguage): List<Virtu
     ModuleUtilCore.collectModulesDependsOn(this, dependentModules)
 
     return dependentModules
-        .flatMap { it.suitableTestSourceFolders(codegenLanguage) }
-        .mapNotNull { it.file }
+        .flatMap { it.suitableTestSourceFolders() }
+        .mapNotNull { it.testSourceRoot }
 }
 
-private fun Module.suitableTestSourceFolders(codegenLanguage: CodegenLanguage): List<SourceFolder> {
+private val SourceFolder.testSourceRoot:TestSourceRoot?
+    get() {
+        val file = file
+        val expectedLanguage = expectedLanguageForTests
+        if (file != null && expectedLanguage != null)
+            return TestSourceRoot(file, expectedLanguage)
+        return null
+    }
+
+private fun Module.suitableTestSourceFolders(): List<SourceFolder> {
     val sourceFolders = ModuleRootManager.getInstance(this)
         .contentEntries
         .flatMap { it.sourceFolders.toList() }
         .filterNotNull()
 
     return sourceFolders
         .filterNot { it.isForGeneratedSources() }
-        .filter { it.rootType == codegenLanguage.testRootType() }
-        // Heuristics: User is more likely to choose the shorter path
-        .sortedBy { it.url.length }
+        .filter { it.isTestSource }
 }
+
 private val GRADLE_SYSTEM_ID = ProjectSystemId("GRADLE")
 
 val Project.isBuildWithGradle get() =
@@ -157,19 +164,20 @@ val Project.isBuildWithGradle get() =
          }
 
 private const val dedicatedTestSourceRootName = "utbot_tests"
-fun Module.addDedicatedTestRoot(testSourceRoots: MutableList<VirtualFile>): VirtualFile? {
+
+fun Module.addDedicatedTestRoot(testSourceRoots: MutableList<TestSourceRoot>, language: CodegenLanguage): VirtualFile? {
     // Don't suggest new test source roots for Gradle project where 'unexpected' test roots won't work
     if (project.isBuildWithGradle) return null
     // Dedicated test root already exists
-    if (testSourceRoots.any { file -> file.name == dedicatedTestSourceRootName }) return null
+    if (testSourceRoots.any { root -> root.dir.name == dedicatedTestSourceRootName }) return null
 
     val moduleInstance = ModuleRootManager.getInstance(this)
     val testFolder = moduleInstance.contentEntries.flatMap { it.sourceFolders.toList() }
         .firstOrNull { it.rootType in testSourceRootTypes }
     (testFolder?.let { testFolder.file?.parent }
         ?: testFolder?.contentEntry?.file ?: this.guessModuleDir())?.let {
         val file = FakeVirtualFile(it, dedicatedTestSourceRootName)
-        testSourceRoots.add(file)
+        testSourceRoots.add(TestSourceRoot(file, language))
         // We return "true" IFF it's case of not yet created fake directory
         return if (VfsUtil.findRelativeFile(it, dedicatedTestSourceRootName) == null) file else null
     }
@@ -275,3 +283,20 @@ private fun jdkVersionBy(sdk: Sdk?): JavaSdkVersion {
     }
     return jdkVersion
 }
+
+private val SourceFolder.expectedLanguageForTests: CodegenLanguage?
+    get() {
+        // unfortunately, Gradle creates Kotlin test source root with Java source root type, so type is misleading,
+        // and we should try looking for name first
+        if (file?.name == "kotlin")
+            return CodegenLanguage.KOTLIN
+
+        if (file?.name == "java")
+            return CodegenLanguage.JAVA
+
+        return when (rootType) {
+            CodegenLanguage.KOTLIN.testRootType() -> CodegenLanguage.KOTLIN
+            CodegenLanguage.JAVA.testRootType() -> CodegenLanguage.JAVA
+            else -> null
+        }
+    }