Add support for nested classes to UI #649

utbot-core/src/main/kotlin/org/utbot/common/KClassUtil.kt

@@ -2,7 +2,7 @@ package org.utbot.common
 
 import java.lang.reflect.InvocationTargetException
 import java.lang.reflect.Method
-
+import kotlin.reflect.KClass
 
 val Class<*>.packageName: String get() = `package`?.name?:""
 
@@ -11,3 +11,6 @@ fun Method.invokeCatching(obj: Any?, args: List<Any?>) = try {
 } catch (e: InvocationTargetException) {
     Result.failure<Nothing>(e.targetException)
 }
+
+val KClass<*>.allNestedClasses: List<KClass<*>>
+    get() = listOf(this) + nestedClasses.flatMap { it.allNestedClasses }

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

@@ -44,6 +44,7 @@ import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.scripting.resolve.classId
 import org.utbot.common.HTML_LINE_SEPARATOR
 import org.utbot.common.PathUtil.toHtmlLinkTag
+import org.utbot.common.allNestedClasses
 import org.utbot.common.appendHtmlLine
 import org.utbot.framework.codegen.Import
 import org.utbot.framework.codegen.ParametrizedTestSource
@@ -70,6 +71,7 @@ import org.utbot.intellij.plugin.ui.WarningTestsReportNotifier
 import org.utbot.intellij.plugin.ui.utils.getOrCreateSarifReportsPath
 import org.utbot.intellij.plugin.ui.utils.showErrorDialogLater
 import org.utbot.intellij.plugin.util.RunConfigurationHelper
+import org.utbot.intellij.plugin.util.extractClassMethodsIncludingNested
 import org.utbot.intellij.plugin.util.signature
 import org.utbot.sarif.SarifReport
 import java.nio.file.Path
@@ -82,7 +84,11 @@ import org.utbot.intellij.plugin.util.IntelliJApiHelper.run
 
 object CodeGenerationController {
 
-    fun generateTests(model: GenerateTestsModel, testSetsByClass: Map<PsiClass, List<UtMethodTestSet>>) {
+    fun generateTests(
+        model: GenerateTestsModel,
+        testSetsByClass: Map<PsiClass, List<UtMethodTestSet>>,
+        psi2KClass: Map<PsiClass, KClass<*>>
+    ) {
         val baseTestDirectory = model.testSourceRoot?.toPsiDirectory(model.project)
             ?: return
         val allTestPackages = getPackageDirectories(baseTestDirectory)
@@ -98,9 +104,10 @@ object CodeGenerationController {
                 val testDirectory = allTestPackages[classPackageName] ?: baseTestDirectory
                 val testClass = createTestClass(srcClass, testDirectory, model) ?: continue
                 val file = testClass.containingFile
+                val cut = psi2KClass[srcClass] ?: error("Didn't find KClass instance for class ${srcClass.name}")
                 runWriteCommandAction(model.project, "Generate tests with UtBot", null, {
                     try {
-                        generateCodeAndReport(srcClass, testClass, file, testSets, model, latch, reports)
+                        generateCodeAndReport(srcClass, cut, testClass, file, testSets, model, latch, reports)
                         testFiles.add(file)
                     } catch (e: IncorrectOperationException) {
                         showCreatingClassError(model.project, createTestClassName(srcClass))
@@ -240,15 +247,15 @@ object CodeGenerationController {
 
     private fun generateCodeAndReport(
         srcClass: PsiClass,
+        classUnderTest: KClass<*>,
         testClass: PsiClass,
         file: PsiFile,
         testSets: List<UtMethodTestSet>,
         model: GenerateTestsModel,
         reportsCountDown: CountDownLatch,
         reports: MutableList<TestsGenerationReport>,
     ) {
-        val classUnderTest = testSets.first().method.clazz
-        val classMethods = TestIntegrationUtils.extractClassMethods(srcClass, false)
+        val classMethods = srcClass.extractClassMethodsIncludingNested(false)
         val paramNames = DumbService.getInstance(model.project)
             .runReadActionInSmartMode(Computable { findMethodParamNames(classUnderTest, classMethods) })
 
@@ -339,7 +346,7 @@ object CodeGenerationController {
 
     private fun findMethodParamNames(clazz: KClass<*>, methods: List<MemberInfo>): Map<ExecutableId, List<String>> {
         val bySignature = methods.associate { it.signature() to it.paramNames() }
-        return clazz.functions
+        return clazz.allNestedClasses.flatMap { it.functions }
             .mapNotNull { method -> bySignature[method.signature()]?.let { params -> method.executableId to params } }
             .toMap()
     }

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

@@ -19,15 +19,14 @@ import com.intellij.openapi.ui.Messages
 import com.intellij.openapi.util.Computable
 import com.intellij.openapi.util.text.StringUtil
 import com.intellij.psi.PsiClass
-import com.intellij.psi.SyntheticElement
+import com.intellij.psi.PsiMethod
 import com.intellij.refactoring.util.classMembers.MemberInfo
-import com.intellij.testIntegration.TestIntegrationUtils
 import com.intellij.util.concurrency.AppExecutorUtil
 import mu.KotlinLogging
 import org.jetbrains.kotlin.idea.util.module
 import org.utbot.analytics.EngineAnalyticsContext
 import org.utbot.analytics.Predictors
-import org.utbot.common.filterWhen
+import org.utbot.common.allNestedClasses
 import org.utbot.engine.util.mockListeners.ForceMockListener
 import org.utbot.framework.plugin.services.JdkInfoService
 import org.utbot.framework.UtSettings
@@ -56,11 +55,11 @@ import java.util.concurrent.TimeUnit
 import org.utbot.engine.util.mockListeners.ForceStaticMockListener
 import org.utbot.framework.PathSelectorType
 import org.utbot.framework.plugin.services.WorkingDirService
+import org.utbot.intellij.plugin.models.packageName
 import org.utbot.intellij.plugin.settings.Settings
+import org.utbot.intellij.plugin.util.extractClassMethodsIncludingNested
 import org.utbot.intellij.plugin.ui.utils.isBuildWithGradle
-import org.utbot.intellij.plugin.ui.utils.suitableTestSourceRoots
 import org.utbot.intellij.plugin.util.PluginWorkingDirProvider
-import org.utbot.intellij.plugin.util.isAbstract
 import kotlin.reflect.KClass
 import kotlin.reflect.full.functions
 
@@ -71,16 +70,18 @@ object UtTestsDialogProcessor {
     fun createDialogAndGenerateTests(
         project: Project,
         srcClasses: Set<PsiClass>,
+        extractMembersFromSrcClasses: Boolean,
         focusedMethod: MemberInfo?,
     ) {
-        createDialog(project, srcClasses, focusedMethod)?.let {
+        createDialog(project, srcClasses, extractMembersFromSrcClasses, focusedMethod)?.let {
             if (it.showAndGet()) createTests(project, it.model)
         }
     }
 
     private fun createDialog(
         project: Project,
         srcClasses: Set<PsiClass>,
+        extractMembersFromSrcClasses: Boolean,
         focusedMethod: MemberInfo?,
     ): GenerateTestsDialogWindow? {
         val srcModule = findSrcModule(srcClasses)
@@ -105,6 +106,7 @@ object UtTestsDialogProcessor {
                 srcModule,
                 testModules,
                 srcClasses,
+                extractMembersFromSrcClasses,
                 if (focusedMethod != null) setOf(focusedMethod) else null,
                 UtSettings.utBotGenerationTimeoutInMillis,
             )
@@ -145,6 +147,7 @@ object UtTestsDialogProcessor {
                             val context = UtContext(classLoader)
 
                             val testSetsByClass = mutableMapOf<PsiClass, List<UtMethodTestSet>>()
+                            val psi2KClass = mutableMapOf<PsiClass, KClass<*>>()
                             var processedClasses = 0
                             val totalClasses = model.srcClasses.size
 
@@ -159,18 +162,23 @@ object UtTestsDialogProcessor {
 
                             for (srcClass in model.srcClasses) {
                                 val methods = ReadAction.nonBlocking<List<UtMethod<*>>> {
-                                    val clazz = classLoader.loadClass(srcClass.qualifiedName).kotlin
-                                    val srcMethods =
-                                        model.selectedMethods?.toList() ?: TestIntegrationUtils.extractClassMethods(
-                                            srcClass,
-                                            false
-                                        )
-                                            .filterWhen(UtSettings.skipTestGenerationForSyntheticMethods) {
-                                                it.member !is SyntheticElement
-                                            }
-                                            .filterNot { it.isAbstract }
+                                    val canonicalName = srcClass.canonicalName
+                                    val clazz = classLoader.loadClass(canonicalName).kotlin
+                                    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()
+                                        chosenMethods + chosenNestedClasses.flatMap {
+                                            it.extractClassMethodsIncludingNested(false)
+                                        }
+                                    } else {
+                                        srcClass.extractClassMethodsIncludingNested(false)
+                                    }
                                     DumbService.getInstance(project).runReadActionInSmartMode(Computable {
-                                        findMethodsInClassMatchingSelected(clazz, srcMethods)
+                                        clazz.allNestedClasses.flatMap {
+                                            findMethodsInClassMatchingSelected(it, srcMethods)
+                                        }
                                     })
                                 }.executeSynchronously()
 
@@ -264,7 +272,7 @@ object UtTestsDialogProcessor {
 
                             invokeLater {
                                 withUtContext(context) {
-                                    generateTests(model, testSetsByClass)
+                                    generateTests(model, testSetsByClass, psi2KClass)
                                 }
                             }
                         }
@@ -273,6 +281,19 @@ object UtTestsDialogProcessor {
             }
     }
 
+    private val PsiClass.canonicalName: String
+        get() {
+            return if (packageName.isEmpty()) {
+                qualifiedName?.replace(".", "$") ?: ""
+            } else {
+                val name = qualifiedName
+                    ?.substringAfter("$packageName.")
+                    ?.replace(".", "$")
+                    ?: ""
+                "$packageName.$name"
+            }
+        }
+
     /**
      * Configures utbot-analytics models for the better path selection.
      *

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

@@ -26,8 +26,9 @@ data class GenerateTestsModel(
     val srcModule: Module,
     val potentialTestModules: List<Module>,
     var srcClasses: Set<PsiClass>,
-    var selectedMethods: Set<MemberInfo>?,
-    var timeout:Long,
+    val extractMembersFromSrcClasses: Boolean,
+    var selectedMembers: Set<MemberInfo>?, // TODO: maybe we should make it not nullable?
+    var timeout: Long,
     var generateWarningsForStaticMocking: Boolean = false,
     var fuzzingValue: Double = 0.05
 ) {
@@ -36,6 +37,7 @@ data class GenerateTestsModel(
     var testModule: Module = potentialTestModules.firstOrNull() ?: error("Empty list of test modules in model")
 
     var testSourceRoot: VirtualFile? = null
+
     fun setSourceRootAndFindTestModule(newTestSourceRoot: VirtualFile?) {
         requireNotNull(newTestSourceRoot)
         testSourceRoot = newTestSourceRoot

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

@@ -39,14 +39,11 @@ import com.intellij.openapi.vfs.newvfs.impl.FakeVirtualFile
 import com.intellij.openapi.wm.ToolWindowManager
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiManager
-import com.intellij.psi.PsiMethod
-import com.intellij.psi.SyntheticElement
 import com.intellij.refactoring.PackageWrapper
 import com.intellij.refactoring.ui.MemberSelectionTable
 import com.intellij.refactoring.ui.PackageNameReferenceEditorCombo
 import com.intellij.refactoring.util.RefactoringUtil
 import com.intellij.refactoring.util.classMembers.MemberInfo
-import com.intellij.testIntegration.TestIntegrationUtils
 import com.intellij.ui.ColoredListCellRenderer
 import com.intellij.ui.ContextHelpLabel
 import com.intellij.ui.HyperlinkLabel
@@ -97,7 +94,6 @@ import org.jetbrains.concurrency.Promise
 import org.jetbrains.concurrency.thenRun
 import org.jetbrains.kotlin.asJava.classes.KtUltraLightClass
 import org.utbot.common.PathUtil.toPath
-import org.utbot.common.filterWhen
 import org.utbot.framework.UtSettings
 import org.utbot.framework.codegen.ForceStaticMocking
 import org.utbot.framework.codegen.Junit4
@@ -136,7 +132,7 @@ import org.utbot.intellij.plugin.ui.utils.parseVersion
 import org.utbot.intellij.plugin.ui.utils.testResourceRootTypes
 import org.utbot.intellij.plugin.ui.utils.testRootType
 import org.utbot.intellij.plugin.util.IntelliJApiHelper
-import org.utbot.intellij.plugin.util.isAbstract
+import org.utbot.intellij.plugin.util.extractFirstLevelMembers
 
 private const val RECENTS_KEY = "org.utbot.recents"
 
@@ -384,17 +380,14 @@ class GenerateTestsDialogWindow(val model: GenerateTestsModel) : DialogWrapper(m
     private fun updateMembersTable() {
         val srcClasses = model.srcClasses
 
-        val items: List<MemberInfo>
-        if (srcClasses.size == 1) {
-            items = TestIntegrationUtils.extractClassMethods(srcClasses.single(), false)
-                .filterWhen(UtSettings.skipTestGenerationForSyntheticMethods) { it.member !is SyntheticElement }
-                .filterNot { it.isAbstract }
-            updateMethodsTable(items)
+        val items = if (model.extractMembersFromSrcClasses) {
+            srcClasses.flatMap { it.extractFirstLevelMembers(false) }
         } else {
-            items = srcClasses.map { MemberInfo(it) }
-            updateClassesTable(items)
+            srcClasses.map { MemberInfo(it) }
         }
 
+        checkMembers(items)
+        membersTable.setMemberInfos(items)
         if (items.isEmpty()) isOKActionEnabled = false
 
         // fix issue with MemberSelectionTable height, set it directly.
@@ -403,28 +396,14 @@ class GenerateTestsDialogWindow(val model: GenerateTestsModel) : DialogWrapper(m
         membersTable.preferredScrollableViewportSize = size(-1, height)
     }
 
-    private fun updateMethodsTable(allMethods: List<MemberInfo>) {
-        val selectedDisplayNames = model.selectedMethods?.map { it.displayName } ?: emptyList()
-        val selectedMethods = if (selectedDisplayNames.isEmpty())
-            allMethods
-            else allMethods.filter { it.displayName in selectedDisplayNames }
+    private fun checkMembers(allMembers: List<MemberInfo>) {
+        val selectedDisplayNames = model.selectedMembers?.map { it.displayName } ?: emptyList()
+        val selectedMembers = allMembers.filter { it.displayName in selectedDisplayNames }
 
-        if (selectedMethods.isEmpty()) {
-            checkMembers(allMethods)
-        } else {
-            checkMembers(selectedMethods)
-        }
-
-        membersTable.setMemberInfos(allMethods)
-    }
-
-    private fun updateClassesTable(srcClasses: List<MemberInfo>) {
-        checkMembers(srcClasses)
-        membersTable.setMemberInfos(srcClasses)
+        val methodsToCheck = selectedMembers.ifEmpty { allMembers }
+        methodsToCheck.forEach { it.isChecked = true }
     }
 
-    private fun checkMembers(members: List<MemberInfo>) = members.forEach { it.isChecked = true }
-
     private fun getTestRoot() : VirtualFile? {
         model.testSourceRoot?.let {
             if (it.isDirectory || it is FakeVirtualFile) return it
@@ -497,12 +476,12 @@ class GenerateTestsDialogWindow(val model: GenerateTestsModel) : DialogWrapper(m
             if (testPackageField.text != SAME_PACKAGE_LABEL) testPackageField.text else ""
 
         val selectedMembers = membersTable.selectedMemberInfos
-        model.srcClasses = selectedMembers
-            .mapNotNull { it.member as? PsiClass ?: it.member.containingClass }
-            .toSet()
-
-        val selectedMethods = selectedMembers.filter { it.member is PsiMethod }.toSet()
-        model.selectedMethods = if (selectedMethods.any()) selectedMethods else null
+        if (!model.extractMembersFromSrcClasses) {
+            model.srcClasses = selectedMembers
+                .mapNotNull { it.member as? PsiClass }
+                .toSet()
+        }
+        model.selectedMembers = selectedMembers.toSet()
 
         model.testFramework = testFrameworks.item
         model.mockStrategy = mockStrategies.item