UnitTestBot · Vassiliy-Kudryashov · Oct 26, 2022 · Oct 7, 2022 · Oct 7, 2022 · Oct 20, 2022
diff --git a/...-intellij/src/main/kotlin/org/utbot/intellij/plugin/generator/CodeGenerationController.kt b/...-intellij/src/main/kotlin/org/utbot/intellij/plugin/generator/CodeGenerationController.kt
@@ -22,6 +22,7 @@ import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.project.DumbService
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Computable
+import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.wm.ToolWindowManager
 import com.intellij.psi.*
 import com.intellij.psi.codeStyle.CodeStyleManager
@@ -73,6 +74,9 @@ import java.nio.file.Path
 import java.util.concurrent.CancellationException
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
+import org.jetbrains.kotlin.idea.util.projectStructure.allModules
+import org.utbot.intellij.plugin.ui.utils.TestSourceRoot
+import org.utbot.intellij.plugin.ui.utils.isBuildWithGradle
 
 object CodeGenerationController {
     private val logger = KotlinLogging.logger {}
@@ -433,6 +437,12 @@ object CodeGenerationController {
         }
     }
 
+    fun GenerateTestsModel.getAllTestSourceRoots() : MutableList<TestSourceRoot> {
+        with(if (project.isBuildWithGradle) project.allModules() else potentialTestModules) {
+            return this.flatMap { it.suitableTestSourceRoots().toList() }.toMutableList()
+        }
+    }
+
     private val CodegenLanguage.utilClassFileName: String
         get() = "$UT_UTILS_CLASS_NAME${this.extension}"
 

diff --git a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/generator/UtTestsDialogProcessor.kt b/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/generator/UtTestsDialogProcessor.kt
@@ -38,9 +38,7 @@ import org.utbot.intellij.plugin.process.EngineProcess
 import org.utbot.intellij.plugin.process.RdTestGenerationResult
 import org.utbot.intellij.plugin.settings.Settings
 import org.utbot.intellij.plugin.ui.GenerateTestsDialogWindow
-import org.utbot.intellij.plugin.ui.utils.isBuildWithGradle
 import org.utbot.intellij.plugin.ui.utils.showErrorDialogLater
-import org.utbot.intellij.plugin.ui.utils.suitableTestSourceRoots
 import org.utbot.intellij.plugin.ui.utils.testModules
 import org.utbot.intellij.plugin.util.*
 import org.utbot.rd.terminateOnException
@@ -49,6 +47,7 @@ import java.nio.file.Path
 import java.nio.file.Paths
 import java.util.concurrent.TimeUnit
 import kotlin.io.path.pathString
+import org.utbot.intellij.plugin.generator.CodeGenerationController.getAllTestSourceRoots
 import org.utbot.framework.plugin.api.util.LockFile
 
 object UtTestsDialogProcessor {
@@ -97,7 +96,16 @@ object UtTestsDialogProcessor {
         // we want to start the child process in the same directory as the test runner
         WorkingDirService.workingDirProvider = PluginWorkingDirProvider(project)
 
-        if (project.isBuildWithGradle && testModules.flatMap { it.suitableTestSourceRoots() }.isEmpty()) {
+        val model = GenerateTestsModel(
+            project,
+            srcModule,
+            testModules,
+            srcClasses,
+            extractMembersFromSrcClasses,
+            focusedMethods,
+            UtSettings.utBotGenerationTimeoutInMillis,
+        )
+        if (model.getAllTestSourceRoots().isEmpty()) {
             val errorMessage = """
                 <html>No test source roots found in the project.<br>
                 Please, <a href="https://www.jetbrains.com/help/idea/testing.html#add-test-root">create or configure</a> at least one test source root.
@@ -106,17 +114,7 @@ object UtTestsDialogProcessor {
             return null
         }
 
-        return GenerateTestsDialogWindow(
-            GenerateTestsModel(
-                project,
-                srcModule,
-                testModules,
-                srcClasses,
-                extractMembersFromSrcClasses,
-                focusedMethods,
-                UtSettings.utBotGenerationTimeoutInMillis,
-            )
-        )
+        return GenerateTestsDialogWindow(model)
     }
 
     private fun createTests(project: Project, model: GenerateTestsModel) {

diff --git a/...rc/main/kotlin/org/utbot/intellij/plugin/ui/components/TestFolderComboWithBrowseButton.kt b/...rc/main/kotlin/org/utbot/intellij/plugin/ui/components/TestFolderComboWithBrowseButton.kt
@@ -8,22 +8,24 @@ import com.intellij.openapi.project.guessProjectDir
 import com.intellij.openapi.ui.ComboBox
 import com.intellij.openapi.ui.ComponentWithBrowseButton
 import com.intellij.openapi.ui.FixedSizeButton
+import com.intellij.openapi.util.text.StringUtil
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.newvfs.impl.FakeVirtualFile
 import com.intellij.ui.ColoredListCellRenderer
 import com.intellij.ui.SimpleTextAttributes
 import com.intellij.util.ArrayUtil
 import com.intellij.util.ui.UIUtil
 import java.io.File
+import java.util.Comparator
 import javax.swing.DefaultComboBoxModel
 import javax.swing.JList
-import org.jetbrains.kotlin.idea.util.projectStructure.allModules
+import org.jetbrains.kotlin.idea.util.rootManager
 import org.utbot.common.PathUtil
+import org.utbot.intellij.plugin.generator.CodeGenerationController.getAllTestSourceRoots
 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
 
 class TestFolderComboWithBrowseButton(private val model: GenerateTestsModel) :
     ComponentWithBrowseButton<ComboBox<Any>>(ComboBox(), null) {
@@ -57,20 +59,42 @@ class TestFolderComboWithBrowseButton(private val model: GenerateTestsModel) :
             }
         }
 
-        val suggestedModules =
-            if (model.project.isBuildWithGradle) model.project.allModules() else model.potentialTestModules
+        var commonModuleSourceDirectory = ""
+        for ((i, sourceRoot) in model.srcModule.rootManager.sourceRoots.withIndex()) {
+            commonModuleSourceDirectory = if (i == 0) {
+                sourceRoot.toNioPath().toString()
+            } else {
+                StringUtil.commonPrefix(commonModuleSourceDirectory, sourceRoot.toNioPath().toString())
+            }
+        }
+        // The first sorting to obtain the best candidate
+        val testRoots = model.getAllTestSourceRoots().distinct().sortedWith(object : Comparator<TestSourceRoot> {
+            override fun compare(o1: TestSourceRoot, o2: TestSourceRoot): Int {
+                // Heuristics: Dirs with language == codegenLanguage should go first
+                val languageOrder = (o1.expectedLanguage == model.codegenLanguage).compareTo(o2.expectedLanguage == model.codegenLanguage)
+                if (languageOrder != 0) return -languageOrder
+                // Heuristics: move root that is 'closer' to module 'common' directory to the first position
+                return -StringUtil.commonPrefixLength(commonModuleSourceDirectory, o1.dir.toNioPath().toString())
+                    .compareTo(StringUtil.commonPrefixLength(commonModuleSourceDirectory, o2.dir.toNioPath().toString()))
+            }
+        }).toMutableList()
+
+        val theBest = if (testRoots.isNotEmpty()) testRoots[0] else null
 
-        val testRoots = suggestedModules.flatMap {
-            it.suitableTestSourceRoots()
-        }.sortedWith(
-            compareByDescending<TestSourceRoot> {
+        // The second sorting to make full list ordered
+        testRoots.sortWith(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
+                // ABC-sorting
+                it.dir.toNioPath()
             }
-        ).toMutableList()
+        )
+        // The best candidate should go first to be pre-selected
+        theBest?.let {
+            testRoots.remove(it)
+            testRoots.add(0, it)
+        }
 
         // this method is blocked for Gradle, where multiple test modules can exist
         model.testModule.addDedicatedTestRoot(testRoots, model.codegenLanguage)