diff --git a/utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonLanguageAssistant.kt b/utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonLanguageAssistant.kt
index 8339d80f2d..34e94caa6c 100644
--- a/utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonLanguageAssistant.kt
+++ b/utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonLanguageAssistant.kt
@@ -44,13 +44,18 @@ object PythonLanguageAssistant : LanguageAssistant() {
     }
 
     private fun getPsiTargets(e: AnActionEvent): Targets? {
-        val editor = e.getData(CommonDataKeys.EDITOR) ?: return null
+        val editor = e.getData(CommonDataKeys.EDITOR)
         val file = e.getData(CommonDataKeys.PSI_FILE) as? PyFile ?: return null
-        val element = findPsiElement(file, editor) ?: return null
 
         if (file.module?.sdk?.sdkType !is PythonSdkType)
             return null
 
+        val element = if (editor != null) {
+            findPsiElement(file, editor) ?: return null
+        } else {
+            e.getData(CommonDataKeys.PSI_ELEMENT) ?: return null
+        }
+
         val containingFunction = getContainingElement<PyFunction>(element)
         val containingClass = getContainingElement<PyClass>(element)
 
@@ -67,7 +72,8 @@ object PythonLanguageAssistant : LanguageAssistant() {
         if (functions.isEmpty())
             return null
 
-        val focusedFunction = if (functions.any { it.name == containingFunction?.name }) containingFunction else null
+        val focusedFunction =
+            if (functions.any { it.name == containingFunction?.name }) containingFunction else null
         return Targets(functions.toSet(), containingClass, focusedFunction, file)
     }
 
diff --git a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/actions/GenerateTestsAction.kt b/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/actions/GenerateTestsAction.kt
index a8c395b385..b06d0c0598 100644
--- a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/actions/GenerateTestsAction.kt
+++ b/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/actions/GenerateTestsAction.kt
@@ -10,6 +10,11 @@ class GenerateTestsAction : AnAction() {
     }
 
     override fun update(e: AnActionEvent) {
-        LanguageAssistant.get(e)?.update(e)
+        val languageAssistant = LanguageAssistant.get(e)
+        if (languageAssistant == null) {
+            e.presentation.isEnabled = false
+        } else {
+            languageAssistant.update(e)
+        }
     }
 }
\ No newline at end of file