diff --git a/src/main/kotlin/org/jetbrains/research/testspark/actions/GenerateTestsUtils.kt b/src/main/kotlin/org/jetbrains/research/testspark/actions/GenerateTestsUtils.kt
index 8a8e112cc..925882cc4 100644
--- a/src/main/kotlin/org/jetbrains/research/testspark/actions/GenerateTestsUtils.kt
+++ b/src/main/kotlin/org/jetbrains/research/testspark/actions/GenerateTestsUtils.kt
@@ -14,6 +14,7 @@ import com.intellij.psi.PsiCodeBlock
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiJavaFile
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiStatement
 import com.intellij.psi.util.PsiTreeUtil