diff --git a/src/main/kotlin/org/jetbrains/research/testspark/display/utils/java/JavaDisplayUtils.kt b/src/main/kotlin/org/jetbrains/research/testspark/display/utils/java/JavaDisplayUtils.kt
index 0320308d7..d068c01f5 100644
--- a/src/main/kotlin/org/jetbrains/research/testspark/display/utils/java/JavaDisplayUtils.kt
+++ b/src/main/kotlin/org/jetbrains/research/testspark/display/utils/java/JavaDisplayUtils.kt
@@ -15,7 +15,6 @@ import com.intellij.psi.PsiElementFactory
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiJavaFile
 import com.intellij.psi.PsiManager
-import com.intellij.refactoring.suggested.startOffset
 import com.intellij.util.containers.stream
 import org.jetbrains.research.testspark.bundles.plugin.PluginLabelsBundle
 import org.jetbrains.research.testspark.data.UIContext
@@ -202,7 +201,7 @@ class JavaDisplayUtils : DisplayUtils {
 
         // insert imports to a code
         PsiDocumentManager.getInstance(project).getDocument(outputFile)!!.insertString(
-            outputFile.importList?.startOffset ?: outputFile.packageStatement?.startOffset ?: 0,
+            outputFile.importList?.startOffsetInParent ?: outputFile.packageStatement?.startOffsetInParent ?: 0,
             uiContext.testGenerationOutput.importsCode.joinToString("\n") + "\n\n",
         )
 
diff --git a/src/main/kotlin/org/jetbrains/research/testspark/display/utils/kotlin/KotlinDisplayUtils.kt b/src/main/kotlin/org/jetbrains/research/testspark/display/utils/kotlin/KotlinDisplayUtils.kt
index f2d61231a..3f22a87ac 100644
--- a/src/main/kotlin/org/jetbrains/research/testspark/display/utils/kotlin/KotlinDisplayUtils.kt
+++ b/src/main/kotlin/org/jetbrains/research/testspark/display/utils/kotlin/KotlinDisplayUtils.kt
@@ -13,12 +13,11 @@ import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiJavaFile
 import com.intellij.psi.PsiManager
-import com.intellij.refactoring.suggested.endOffset
-import com.intellij.refactoring.suggested.startOffset
 import com.intellij.util.containers.stream
 import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.jetbrains.kotlin.psi.psiUtil.pureEndOffset
 import org.jetbrains.research.testspark.bundles.plugin.PluginLabelsBundle
 import org.jetbrains.research.testspark.data.UIContext
 import org.jetbrains.research.testspark.display.utils.ErrorMessageManager
@@ -213,8 +212,8 @@ class KotlinDisplayUtils : DisplayUtils {
         val importsString = uiContext.testGenerationOutput.importsCode.joinToString("\n") + "\n\n"
 
         // Find the insertion offset
-        val insertionOffset = outputFile.importList?.startOffset
-            ?: outputFile.packageDirective?.endOffset
+        val insertionOffset = outputFile.importList?.startOffsetInParent
+            ?: outputFile.packageDirective?.pureEndOffset
             ?: 0
 
         // Insert the imports into the document