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
  }