From 4d46f340c241fa37b5f440e7281f84588f7171a8 Mon Sep 17 00:00:00 2001
From: Ivan Volkov <65076429+volivan239@users.noreply.github.com>
Date: Fri, 7 Oct 2022 12:47:01 +0300
Subject: [PATCH] Get PsiClass' package by declaration in containingFile and
 not by its directory #1081 (#1111)

---
 .../plugin/generator/UtTestsDialogProcessor.kt       |  2 +-
 .../intellij/plugin/models/GenerateTestsModel.kt     | 12 ++++++++++--
 2 files changed, 11 insertions(+), 3 deletions(-)

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 85104b3100..76dd002c2f 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
@@ -271,7 +271,7 @@ object UtTestsDialogProcessor {
                 val name = qualifiedName
                     ?.substringAfter("$packageName.")
                     ?.replace(".", "$")
-                    ?: ""
+                    ?: error("Unable to get canonical name for $this")
                 "$packageName.$name"
             }
         }
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 0aecbc8629..cf1a330978 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
@@ -17,8 +17,9 @@ import com.intellij.openapi.projectRoots.JavaSdkVersion
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.newvfs.impl.FakeVirtualFile
 import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiJavaFile
 import com.intellij.refactoring.util.classMembers.MemberInfo
-import org.jetbrains.kotlin.idea.core.getPackage
+import org.jetbrains.kotlin.psi.KtFile
 import org.utbot.framework.plugin.api.JavaDocCommentStyle
 import org.utbot.framework.util.ConflictTriggers
 import org.utbot.intellij.plugin.ui.utils.jdkVersion
@@ -84,4 +85,11 @@ data class GenerateTestsModel(
         }
 }
 
-val PsiClass.packageName: String get() = this.containingFile.containingDirectory.getPackage()?.qualifiedName ?: ""
\ No newline at end of file
+val PsiClass.packageName: String
+    get() {
+        return when (val currentFile = containingFile) {
+            is PsiJavaFile -> currentFile.packageName
+            is KtFile -> currentFile.packageFqName.asString()
+            else -> error("Can't find package name for $this: it should be located either in Java or Kt file")
+        }
+    }
\ No newline at end of file