Select Kotlin test src root by default when generating for Kotlin

Author: volivan239

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

@@ -18,6 +18,7 @@ import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.newvfs.impl.FakeVirtualFile
 import com.intellij.psi.PsiClass
 import com.intellij.refactoring.util.classMembers.MemberInfo
+import org.jetbrains.kotlin.asJava.classes.KtUltraLightClass
 import org.jetbrains.kotlin.idea.core.getPackage
 import org.utbot.framework.plugin.api.JavaDocCommentStyle
 import org.utbot.framework.util.ConflictTriggers
@@ -55,13 +56,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

@@ -96,7 +96,6 @@ import javax.swing.text.DefaultFormatter
 import kotlin.streams.toList
 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
@@ -652,8 +651,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

@@ -18,6 +18,7 @@ import javax.swing.DefaultComboBoxModel
 import javax.swing.JList
 import org.jetbrains.kotlin.idea.util.projectStructure.allModules
 import org.utbot.common.PathUtil
+import org.utbot.framework.plugin.api.CodegenLanguage
 import org.utbot.intellij.plugin.models.GenerateTestsModel
 import org.utbot.intellij.plugin.ui.utils.addDedicatedTestRoot
 import org.utbot.intellij.plugin.ui.utils.isBuildWithGradle
@@ -58,7 +59,17 @@ 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 testRootsByLanguage = CodegenLanguage.allItems.associateWith { language ->
+            suggestedModules.flatMap { module ->
+                module.suitableTestSourceRoots(language)
+            }
+        }
+
+        // testRoots for default codegen language should go before other testRoots (this impacts default-selected test root item)
+        val testRoots = with (testRootsByLanguage.entries) {
+            filter { it.key == model.codegenLanguage } + filter { it.key != model.codegenLanguage }
+        }.flatMap { it.value }.toMutableList()
+
         // this method is blocked for Gradle, where multiple test modules can exist
         model.testModule.addDedicatedTestRoot(testRoots)
 
@@ -95,7 +106,6 @@ class TestFolderComboWithBrowseButton(private val model: GenerateTestsModel) :
         }
 
     private fun configureRootsCombo(testRoots: List<VirtualFile>) {
-        // unfortunately, Gradle creates Kotlin test source root with Java source root type, so type is misleading
         val selectedRoot = testRoots.first()
 
         // do not update model.testModule here, because fake test source root could have been chosen

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

@@ -145,7 +145,8 @@ private fun Module.suitableTestSourceFolders(codegenLanguage: CodegenLanguage):
 
     return sourceFolders
         .filterNot { it.isForGeneratedSources() }
-        .filter { it.rootType == codegenLanguage.testRootType() }
+        .filter { it.isTestSource }
+        .filter { it.expectedLanguageForTests == codegenLanguage }
         // Heuristics: User is more likely to choose the shorter path
         .sortedBy { it.url.length }
 }
@@ -275,3 +276,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
+        }
+    }