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 index 78a8f6209b..05d58a6cb8 100644 --- 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 @@ -72,9 +72,9 @@ object UtTestsDialogProcessor { project: Project, srcClasses: Set, extractMembersFromSrcClasses: Boolean, - focusedMethod: MemberInfo?, + focusedMethods: Set, ) { - 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, extractMembersFromSrcClasses: Boolean, - focusedMethod: MemberInfo?, + focusedMethods: Set, ): 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) } diff --git a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/models/GenerateTestsModel.kt b/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/models/GenerateTestsModel.kt index 5c4928a241..584b74b9ef 100644 --- a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/models/GenerateTestsModel.kt +++ b/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/models/GenerateTestsModel.kt @@ -27,7 +27,7 @@ data class GenerateTestsModel( val potentialTestModules: List, var srcClasses: Set, val extractMembersFromSrcClasses: Boolean, - var selectedMembers: Set?, // TODO: maybe we should make it not nullable? + var selectedMembers: Set, var timeout: Long, var generateWarningsForStaticMocking: Boolean = false, var fuzzingValue: Double = 0.05 diff --git a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/GenerateTestsDialogWindow.kt b/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/GenerateTestsDialogWindow.kt index 7460f0763c..d4417bff1a 100644 --- a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/GenerateTestsDialogWindow.kt +++ b/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) { - 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 } diff --git a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/actions/GenerateTestsAction.kt b/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/actions/GenerateTestsAction.kt index ee7201bd5c..e1f2951261 100644 --- a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/actions/GenerateTestsAction.kt +++ b/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, MemberInfo?, Boolean>? { + private fun getPsiTargets(e: AnActionEvent): Triple, Set, 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() - var selectedMethod: MemberInfo? = null + val selectedMethods = mutableSetOf() 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 }