Implement search for all @configuration classes in Spring project #1921 (#1934)

Author: IlyaMuravjov

utbot-ui-commons/src/main/kotlin/org/utbot/intellij/plugin/models/BaseTestModel.kt

@@ -1,15 +1,25 @@
 package org.utbot.intellij.plugin.models
 
 import com.intellij.openapi.module.Module
+import com.intellij.openapi.module.ModuleManager
 import com.intellij.openapi.module.ModuleUtil
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.roots.TestModuleProperties
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.newvfs.impl.FakeVirtualFile
+import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.PsiClass
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.ProjectScope
+import com.intellij.psi.search.searches.AnnotatedElementsSearch
 import org.jetbrains.kotlin.idea.core.getPackage
+import org.jetbrains.kotlin.idea.search.allScope
 import org.jetbrains.kotlin.idea.util.projectStructure.allModules
+import org.jetbrains.kotlin.idea.util.rootManager
+import org.jetbrains.kotlin.idea.util.sourceRoot
 import org.utbot.framework.plugin.api.CodegenLanguage
 import org.utbot.intellij.plugin.ui.utils.ITestSourceRoot
+import org.utbot.intellij.plugin.ui.utils.getSortedTestRoots
 import org.utbot.intellij.plugin.ui.utils.isBuildWithGradle
 import org.utbot.intellij.plugin.ui.utils.suitableTestSourceRoots
 
@@ -56,6 +66,51 @@ open class BaseTestsModel(
         }
     }
 
+    fun getSortedTestRoots(): MutableList<ITestSourceRoot> = getSortedTestRoots(
+        getAllTestSourceRoots(),
+        sourceRootHistory,
+        srcModule.rootManager.sourceRoots.map { file: VirtualFile -> file.toNioPath().toString() },
+        codegenLanguage
+    )
+
+    /**
+     * Searches configuration classes in Spring application.
+     *
+     * Classes are selected and sorted in the following order:
+     *   - Classes marked with `@TestConfiguration` annotation
+     *   - Classes marked with `@Configuration` annotation
+     *      - firstly, from test source roots (in the order provided by [getSortedTestRoots])
+     *      - after that, from source roots
+     */
+    fun getSortedSpringConfigurationClasses(): List<PsiClass> {
+        val testRootToIndex = getSortedTestRoots().withIndex().associate { (i, root) -> root.dir to i }
+
+        // Not using `srcModule.testModules(project)` here because it returns
+        // test modules for dependent modules if no test roots are found in the source module itself.
+        // We don't want to search configurations there because they seem useless.
+        val testModules = ModuleManager.getInstance(project)
+            .modules
+            .filter { module -> TestModuleProperties.getInstance(module).productionModule == srcModule }
+
+        val searchScope = testModules.fold(GlobalSearchScope.moduleScope(srcModule)) { accScope, module ->
+            accScope.union(GlobalSearchScope.moduleScope(module))
+        }
+
+        val annotationClasses =  listOf(
+            "org.springframework.boot.test.context.TestConfiguration",
+            "org.springframework.context.annotation.Configuration"
+        ).mapNotNull {
+            JavaPsiFacade.getInstance(project).findClass(it, project.allScope())
+        }
+
+        return annotationClasses.flatMap { annotation ->
+            AnnotatedElementsSearch
+                .searchPsiClasses(annotation, searchScope)
+                .findAll()
+                .sortedBy { testRootToIndex[it.containingFile.sourceRoot] ?: Int.MAX_VALUE }
+        }
+    }
+
     fun updateSourceRootHistory(path: String) {
         sourceRootHistory.apply {
             remove(path)//Remove existing entry if any

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

@@ -17,12 +17,10 @@ import com.intellij.util.ui.UIUtil
 import java.io.File
 import javax.swing.DefaultComboBoxModel
 import javax.swing.JList
-import org.jetbrains.kotlin.idea.util.rootManager
 import org.utbot.common.PathUtil
 import org.utbot.intellij.plugin.models.BaseTestsModel
 import org.utbot.intellij.plugin.ui.utils.ITestSourceRoot
 import org.utbot.intellij.plugin.ui.utils.addDedicatedTestRoot
-import org.utbot.intellij.plugin.ui.utils.getSortedTestRoots
 import org.utbot.intellij.plugin.ui.utils.isBuildWithGradle
 
 private const val SET_TEST_FOLDER = "set test folder"
@@ -58,12 +56,7 @@ class TestFolderComboWithBrowseButton(private val model: BaseTestsModel) :
             }
         }
 
-        val testRoots = getSortedTestRoots(
-            model.getAllTestSourceRoots(),
-            model.sourceRootHistory,
-            model.srcModule.rootManager.sourceRoots.map { file: VirtualFile -> file.toNioPath().toString() },
-            model.codegenLanguage
-        )
+        val testRoots = model.getSortedTestRoots()
 
         // this method is blocked for Gradle, where multiple test modules can exist
         model.testModule.addDedicatedTestRoot(testRoots, model.codegenLanguage)