@@ -66,8 +66,7 @@ class GenerateTestsAction : AnAction(), UpdateInBackground {
66
66
67
67
val psiElementHandler = PsiElementHandler .makePsiElementHandler(file)
68
68
69
- // When in Kotlin file, we should propose top-level functions for testing
70
- if (psiElementHandler.isCreateTestActionAvailable(element) || file is KtFile ) {
69
+ if (psiElementHandler.isCreateTestActionAvailable(element)) {
71
70
val srcClass = psiElementHandler.containingClass(element) ? : return null
72
71
val srcSourceRoot = srcClass.getSourceRoot() ? : return null
73
72
val srcMembers = srcClass.extractFirstLevelMembers(false )
@@ -221,7 +220,7 @@ class GenerateTestsAction : AnAction(), UpdateInBackground {
221
220
}
222
221
223
222
private fun getAllClasses (directory : PsiDirectory ): Set <PsiClass > {
224
- val allClasses = directory.files.flatMap { getClassesFromFile(it) }.toMutableSet()
223
+ val allClasses = directory.files.flatMap { PsiElementHandler .makePsiElementHandler(it). getClassesFromFile(it) }.toMutableSet()
225
224
for (subDir in directory.subdirectories) allClasses + = getAllClasses(subDir)
226
225
return allClasses
227
226
}
@@ -234,16 +233,10 @@ class GenerateTestsAction : AnAction(), UpdateInBackground {
234
233
if (! dirsArePackages) {
235
234
return emptySet()
236
235
}
237
- val allClasses = psiFiles.flatMap { getClassesFromFile(it) }.toMutableSet()
236
+ val allClasses = psiFiles.flatMap { PsiElementHandler .makePsiElementHandler(it). getClassesFromFile(it) }.toMutableSet()
238
237
allClasses.addAll(psiFiles.mapNotNull { (it as ? KtFile )?.findFacadeClass() })
239
238
for (psiDir in psiDirectories) allClasses + = getAllClasses(psiDir)
240
239
241
240
return allClasses
242
241
}
243
-
244
- private fun getClassesFromFile (psiFile : PsiFile ): List <PsiClass > {
245
- val psiElementHandler = PsiElementHandler .makePsiElementHandler(psiFile)
246
- return PsiTreeUtil .getChildrenOfTypeAsList(psiFile, psiElementHandler.classClass)
247
- .map { psiElementHandler.toPsi(it, PsiClass ::class .java) }
248
- }
249
242
}
0 commit comments