From 490fc77f769cc8da3f26b6cce33171c2b396d4a3 Mon Sep 17 00:00:00 2001
From: Andrey Tarbeev <github@tarbeev.su>
Date: Tue, 29 Nov 2022 13:35:58 +0300
Subject: [PATCH] Fix menu item "Generate tests ..." for multiple selected kt
 classes

---
 .../kotlin/org/utbot/intellij/plugin/language/JavaLanguage.kt   | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/language/JavaLanguage.kt b/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/language/JavaLanguage.kt
index 0895eb73cf..f2fa701e19 100644
--- a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/language/JavaLanguage.kt
+++ b/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/language/JavaLanguage.kt
@@ -24,6 +24,7 @@ import org.utbot.intellij.plugin.util.extractFirstLevelMembers
 import org.utbot.intellij.plugin.util.isVisible
 import java.util.*
 import org.jetbrains.kotlin.j2k.getContainingClass
+import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.utbot.framework.plugin.api.util.LockFile
 import org.utbot.intellij.plugin.models.packageName
@@ -113,6 +114,7 @@ object JvmLanguageAssistant : LanguageAssistant() {
                     when(it) {
                         is PsiFileSystemItem  -> srcClasses += getAllClasses(project, arrayOf(it.virtualFile))
                         is PsiClass -> srcClasses.add(it)
+                        is KtClass -> srcClasses += getClassesFromFile(it.containingKtFile)
                         is PsiElement -> {
                             srcClasses.addIfNotNull(it.getContainingClass())
                             if (it is PsiMethod) {