UI action is disabled for several selected classes or packages #882 (#896)

Allow multi-selection for packages and methods

Author: Vassiliy-Kudryashov

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

@@ -72,9 +72,9 @@ object UtTestsDialogProcessor {
         project: Project,
         srcClasses: Set<PsiClass>,
         extractMembersFromSrcClasses: Boolean,
-        focusedMethod: MemberInfo?,
+        focusedMethods: Set<MemberInfo>,
     ) {
-        createDialog(project, srcClasses, extractMembersFromSrcClasses, focusedMethod)?.let {
+        createDialog(project, srcClasses, extractMembersFromSrcClasses, focusedMethods)?.let {
             if (it.showAndGet()) createTests(project, it.model)
         }
     }
@@ -83,7 +83,7 @@ object UtTestsDialogProcessor {
         project: Project,
         srcClasses: Set<PsiClass>,
         extractMembersFromSrcClasses: Boolean,
-        focusedMethod: MemberInfo?,
+        focusedMethods: Set<MemberInfo>,
     ): GenerateTestsDialogWindow? {
         val srcModule = findSrcModule(srcClasses)
         val testModules = srcModule.testModules(project)
@@ -108,7 +108,7 @@ object UtTestsDialogProcessor {
                 testModules,
                 srcClasses,
                 extractMembersFromSrcClasses,
-                if (focusedMethod != null) setOf(focusedMethod) else null,
+                focusedMethods,
                 UtSettings.utBotGenerationTimeoutInMillis,
             )
         )
@@ -168,8 +168,8 @@ object UtTestsDialogProcessor {
                                     psi2KClass[srcClass] = clazz
 
                                     val srcMethods = if (model.extractMembersFromSrcClasses) {
-                                        val chosenMethods = model.selectedMembers?.filter { it.member is PsiMethod } ?: listOf()
-                                        val chosenNestedClasses = model.selectedMembers?.mapNotNull { it.member as? PsiClass } ?: listOf()
+                                        val chosenMethods = model.selectedMembers.filter { it.member is PsiMethod }
+                                        val chosenNestedClasses = model.selectedMembers.mapNotNull { it.member as? PsiClass }
                                         chosenMethods + chosenNestedClasses.flatMap {
                                             it.extractClassMethodsIncludingNested(false)
                                         }

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

@@ -27,7 +27,7 @@ data class GenerateTestsModel(
     val potentialTestModules: List<Module>,
     var srcClasses: Set<PsiClass>,
     val extractMembersFromSrcClasses: Boolean,
-    var selectedMembers: Set<MemberInfo>?, // TODO: maybe we should make it not nullable?
+    var selectedMembers: Set<MemberInfo>,
     var timeout: Long,
     var generateWarningsForStaticMocking: Boolean = false,
     var fuzzingValue: Double = 0.05

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

@@ -401,7 +401,7 @@ class GenerateTestsDialogWindow(val model: GenerateTestsModel) : DialogWrapper(m
     }
 
     private fun checkMembers(allMembers: List<MemberInfo>) {
-        val selectedDisplayNames = model.selectedMembers?.map { it.displayName } ?: emptyList()
+        val selectedDisplayNames = model.selectedMembers.map { it.displayName }
         val selectedMembers = allMembers.filter { it.displayName in selectedDisplayNames }
 
         val methodsToCheck = selectedMembers.ifEmpty { allMembers }

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

@@ -5,6 +5,7 @@ import org.utbot.intellij.plugin.ui.utils.PsiElementHandler
 import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.CommonDataKeys
+import com.intellij.openapi.actionSystem.PlatformDataKeys
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.module.ModuleUtil
 import com.intellij.openapi.project.Project
@@ -19,19 +20,21 @@ import org.jetbrains.kotlin.idea.core.util.toPsiDirectory
 import org.jetbrains.kotlin.idea.core.util.toPsiFile
 import org.utbot.intellij.plugin.util.extractFirstLevelMembers
 import java.util.*
+import org.jetbrains.kotlin.j2k.getContainingClass
+import org.jetbrains.kotlin.utils.addIfNotNull
 
 class GenerateTestsAction : AnAction() {
     override fun actionPerformed(e: AnActionEvent) {
         val project = e.project ?: return
-        val (srcClasses, focusedMethod, extractMembersFromSrcClasses) = getPsiTargets(e) ?: return
-        UtTestsDialogProcessor.createDialogAndGenerateTests(project, srcClasses, extractMembersFromSrcClasses, focusedMethod)
+        val (srcClasses, focusedMethods, extractMembersFromSrcClasses) = getPsiTargets(e) ?: return
+        UtTestsDialogProcessor.createDialogAndGenerateTests(project, srcClasses, extractMembersFromSrcClasses, focusedMethods)
     }
 
     override fun update(e: AnActionEvent) {
         e.presentation.isEnabled = getPsiTargets(e) != null
     }
 
-    private fun getPsiTargets(e: AnActionEvent): Triple<Set<PsiClass>, MemberInfo?, Boolean>? {
+    private fun getPsiTargets(e: AnActionEvent): Triple<Set<PsiClass>, Set<MemberInfo>, Boolean>? {
         val project = e.project ?: return null
         val editor = e.getData(CommonDataKeys.EDITOR)
         if (editor != null) {
@@ -56,19 +59,19 @@ class GenerateTestsAction : AnAction() {
                     return null
                 }
 
-                return Triple(setOf(srcClass), focusedMethod, true)
+                return Triple(setOf(srcClass), if (focusedMethod != null) setOf(focusedMethod) else emptySet(), true)
             }
         } else {
             // The action is being called from 'Project' tool window 
             val srcClasses = mutableSetOf<PsiClass>()
-            var selectedMethod: MemberInfo? = null
+            val selectedMethods = mutableSetOf<MemberInfo>()
             var extractMembersFromSrcClasses = false
-            val element = e.getData(CommonDataKeys.PSI_ELEMENT) ?: return null
+            val element = e.getData(CommonDataKeys.PSI_ELEMENT)
             if (element is PsiFileSystemItem) {
                 e.getData(CommonDataKeys.VIRTUAL_FILE_ARRAY)?.let {
                     srcClasses += getAllClasses(project, it)
                 }
-            } else {
+            } else if (element is PsiElement){
                 val file = element.containingFile ?: return null
                 val psiElementHandler = PsiElementHandler.makePsiElementHandler(file)
 
@@ -81,11 +84,29 @@ class GenerateTestsAction : AnAction() {
                     }
 
                     if (element is PsiMethod) {
-                        selectedMethod = MemberInfo(element)
+                        selectedMethods.add(MemberInfo(element))
+                    }
+                }
+            } else {
+                val someSelection = e.getData(PlatformDataKeys.SELECTED_ITEMS)?: return null
+                someSelection.forEach {
+                    when(it) {
+                        is PsiFileSystemItem  -> srcClasses += getAllClasses(project, arrayOf(it.virtualFile))
+                        is PsiClass -> srcClasses.add(it)
+                        is PsiElement -> {
+                            srcClasses.addIfNotNull(it.getContainingClass())
+                            if (it is PsiMethod) {
+                                selectedMethods.add(MemberInfo(it))
+                                extractMembersFromSrcClasses = true
+                            }
+                        }
                     }
                 }
             }
             srcClasses.removeIf { it.isInterface }
+            if (srcClasses.size > 1) {
+                extractMembersFromSrcClasses = false
+            }
             var commonSourceRoot = null as VirtualFile?
             for (srcClass in srcClasses) {
                 if (commonSourceRoot == null) {
@@ -100,7 +121,7 @@ class GenerateTestsAction : AnAction() {
                      .filter { folder -> !folder.rootType.isForTests && folder.file == commonSourceRoot}
                      .findAny().isPresent ) return null
 
-            return Triple(srcClasses.toSet(), selectedMethod, extractMembersFromSrcClasses)
+            return Triple(srcClasses.toSet(), selectedMethods.toSet(), extractMembersFromSrcClasses)
         }
         return null
     }