Move testsourceroot-choosing heuristics from ModuleUtils to TestFolderCombo

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/ui/components/TestFolderComboWithBrowseButton.kt

@@ -18,8 +18,8 @@ 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.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
@@ -59,19 +59,20 @@ class TestFolderComboWithBrowseButton(private val model: GenerateTestsModel) :
         val suggestedModules =
             if (model.project.isBuildWithGradle) model.project.allModules() else model.potentialTestModules
 
-        val testRootsByLanguage = CodegenLanguage.allItems.associateWith { language ->
-            suggestedModules.flatMap { module ->
-                module.suitableTestSourceRoots(language)
+        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
             }
-        }
-
-        // 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()
+        ).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)
@@ -105,12 +106,12 @@ class TestFolderComboWithBrowseButton(private val model: GenerateTestsModel) :
             files.singleOrNull()
         }
 
-    private fun configureRootsCombo(testRoots: List<VirtualFile>) {
+    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,11 +132,20 @@ 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() }
@@ -146,10 +154,8 @@ private fun Module.suitableTestSourceFolders(codegenLanguage: CodegenLanguage):
     return sourceFolders
         .filterNot { it.isForGeneratedSources() }
         .filter { it.isTestSource }
-        .filter { it.expectedLanguageForTests == codegenLanguage }
-        // Heuristics: User is more likely to choose the shorter path
-        .sortedBy { it.url.length }
 }
+
 private val GRADLE_SYSTEM_ID = ProjectSystemId("GRADLE")
 
 val Project.isBuildWithGradle get() =
@@ -158,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
     }