1+ package org.utbot.intellij.plugin.ui.actions
2+
3+ import org.utbot.intellij.plugin.ui.UtTestsDialogProcessor
4+ import org.utbot.intellij.plugin.ui.utils.KotlinPsiElementHandler
5+ import org.utbot.intellij.plugin.ui.utils.PsiElementHandler
6+ import com.intellij.openapi.actionSystem.AnAction
7+ import com.intellij.openapi.actionSystem.AnActionEvent
8+ import com.intellij.openapi.actionSystem.CommonDataKeys
9+ import com.intellij.openapi.editor.Editor
10+ import com.intellij.openapi.module.ModuleUtil
11+ import com.intellij.openapi.project.Project
12+ import com.intellij.openapi.roots.ModuleRootManager
13+ import com.intellij.openapi.roots.ProjectFileIndex
14+ import com.intellij.openapi.vfs.VirtualFile
15+ import com.intellij.psi.*
16+ import com.intellij.psi.util.PsiTreeUtil
17+ import com.intellij.refactoring.util.classMembers.MemberInfo
18+ import com.intellij.testIntegration.TestIntegrationUtils
19+ import org.jetbrains.kotlin.idea.core.getPackage
20+ import org.jetbrains.kotlin.idea.core.util.toPsiDirectory
21+ import org.jetbrains.kotlin.idea.core.util.toPsiFile
22+ import org.jetbrains.kotlin.psi.KtClass
23+ import java.util.*
24+
25+ class GenerateTestsAction : AnAction () {
26+ override fun actionPerformed (e : AnActionEvent ) {
27+ val project = e.project ? : return
28+ val psiTargets = getPsiTargets(e) ? : return
29+ UtTestsDialogProcessor .createDialogAndGenerateTests(project, psiTargets.first, psiTargets.second)
30+ }
31+
32+ override fun update (e : AnActionEvent ) {
33+ e.presentation.isEnabled = getPsiTargets(e) != null
34+ }
35+
36+ private fun getPsiTargets (e : AnActionEvent ): Pair <Set <PsiClass >, MemberInfo?>? {
37+ val project = e.project ? : return null
38+ val editor = e.getData(CommonDataKeys .EDITOR )
39+ if (editor != null ) {
40+ // The action is being called from editor
41+ val file = e.getData(CommonDataKeys .PSI_FILE ) ? : return null
42+ val element = findPsiElement(file, editor) ? : return null
43+
44+ val psiElementHandler = PsiElementHandler .makePsiElementHandler(file)
45+
46+ if (psiElementHandler.isCreateTestActionAvailable(element)) {
47+ val srcClass = psiElementHandler.containingClass(element) ? : return null
48+ val srcMethods = TestIntegrationUtils .extractClassMethods(srcClass, false )
49+ val focusedMethod = focusedMethodOrNull(element, srcMethods, psiElementHandler)
50+ return Pair (setOf (srcClass), focusedMethod)
51+ }
52+ } else {
53+ // The action is being called from 'Project' tool window
54+ val srcClasses = mutableSetOf<PsiClass >()
55+ e.getData(CommonDataKeys .PSI_ELEMENT )?.let {
56+ srcClasses + = getAllClasses(it)
57+ }
58+ e.getData(CommonDataKeys .VIRTUAL_FILE_ARRAY )?.let {
59+ srcClasses + = getAllClasses(project, it)
60+ }
61+ var commonSourceRoot = null as VirtualFile ?
62+ for (srcClass in srcClasses) {
63+ if (commonSourceRoot == null ) {
64+ commonSourceRoot = srcClass.getSourceRoot()? : return null
65+ } else if (commonSourceRoot != srcClass.getSourceRoot()) return null
66+ }
67+ if (commonSourceRoot == null ) return null
68+ val module = ModuleUtil .findModuleForFile(commonSourceRoot, project)? : return null
69+
70+ if (! Arrays .stream(ModuleRootManager .getInstance(module).contentEntries)
71+ .flatMap { entry -> Arrays .stream(entry.sourceFolders) }
72+ .filter { folder -> ! folder.rootType.isForTests && folder.file == commonSourceRoot}
73+ .findAny().isPresent ) return null
74+
75+ return Pair (srcClasses, null )
76+ }
77+ return null
78+ }
79+
80+ private fun PsiElement?.getSourceRoot () : VirtualFile ? {
81+ val project = this ?.project? : return null
82+ val virtualFile = this .containingFile?.originalFile?.virtualFile? : return null
83+ return ProjectFileIndex .getInstance(project).getSourceRootForFile(virtualFile)
84+ }
85+
86+ private fun findPsiElement (file : PsiFile , editor : Editor ): PsiElement ? {
87+ val offset = editor.caretModel.offset
88+ var element = file.findElementAt(offset)
89+ if (element == null && offset == file.textLength) {
90+ element = file.findElementAt(offset - 1 )
91+ }
92+
93+ return element
94+ }
95+
96+ private fun focusedMethodOrNull (element : PsiElement , methods : List <MemberInfo >, psiElementHandler : PsiElementHandler ): MemberInfo ? {
97+ // getParentOfType might return element which does not correspond to the standard Psi hierarchy.
98+ // Thus, make transition to the Psi if it is required.
99+ val currentMethod = PsiTreeUtil .getParentOfType(element, psiElementHandler.methodClass)
100+ ?.let { psiElementHandler.toPsi(it, PsiMethod ::class .java) }
101+
102+ return methods.singleOrNull { it.member == currentMethod }
103+ }
104+
105+ private fun getAllClasses (psiElement : PsiElement ): Set <PsiClass > {
106+ return when (psiElement) {
107+ is KtClass -> setOf (KotlinPsiElementHandler ().toPsi(psiElement, PsiClass ::class .java))
108+ is PsiClass -> setOf (psiElement)
109+ is PsiDirectory -> getAllClasses(psiElement)
110+ else -> emptySet()
111+ }
112+ }
113+
114+ private fun getAllClasses (directory : PsiDirectory ): Set <PsiClass > {
115+ val allClasses = directory.files.flatMap { getClassesFromFile(it) }.toMutableSet()
116+ for (subDir in directory.subdirectories) allClasses + = getAllClasses(subDir)
117+ return allClasses
118+ }
119+ private fun getAllClasses (project : Project , virtualFiles : Array <VirtualFile >): Set <PsiClass > {
120+ val psiFiles = virtualFiles.mapNotNull { it.toPsiFile(project) }
121+ val psiDirectories = virtualFiles.mapNotNull { it.toPsiDirectory(project) }
122+ val dirsArePackages = psiDirectories.all { it.getPackage()?.qualifiedName?.isNotEmpty() == true }
123+
124+ if (! dirsArePackages) {
125+ return emptySet()
126+ }
127+ val allClasses = psiFiles.flatMap { getClassesFromFile(it) }.toMutableSet()
128+ for (psiDir in psiDirectories) allClasses + = getAllClasses(psiDir)
129+
130+ return allClasses
131+ }
132+
133+ private fun getClassesFromFile (psiFile : PsiFile ): List <PsiClass > {
134+ val psiElementHandler = PsiElementHandler .makePsiElementHandler(psiFile)
135+ return PsiTreeUtil .getChildrenOfTypeAsList(psiFile, psiElementHandler.classClass)
136+ .map { psiElementHandler.toPsi(it, PsiClass ::class .java) }
137+ }
138+ }
0 commit comments