diff --git a/idea/idea-gradle/src/org/jetbrains/kotlin/idea/run/kotlinTestClassUtils.kt b/idea/idea-gradle/src/org/jetbrains/kotlin/idea/run/kotlinTestClassUtils.kt
index 60e5843d33d..7c7a8fd4b97 100644
--- a/idea/idea-gradle/src/org/jetbrains/kotlin/idea/run/kotlinTestClassUtils.kt
+++ b/idea/idea-gradle/src/org/jetbrains/kotlin/idea/run/kotlinTestClassUtils.kt
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.run
 import com.intellij.execution.Location
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiMethod
+import org.jetbrains.kotlin.asJava.elements.KtLightElement
 import org.jetbrains.kotlin.idea.caches.lightClasses.KtFakeLightClass
 import org.jetbrains.kotlin.idea.caches.lightClasses.KtFakeLightMethod
 import org.jetbrains.kotlin.psi.KtClassOrObject
@@ -16,13 +17,19 @@ import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 
 internal fun getTestMethodForKotlinTest(location: Location<*>): PsiMethod? {
-    val leaf = location.psiElement
+    val leaf = when (val psi = location.psiElement) {
+        is KtLightElement<*, *> -> psi.kotlinOrigin
+        else -> psi
+    }
     val function = leaf?.getParentOfType<KtNamedFunction>(false) ?: return null
     return KtFakeLightMethod.get(function)
 }
 
 internal fun getTestClassForKotlinTest(location: Location<*>): PsiClass? {
-    val leaf = location.psiElement
-    val owner = leaf?.getParentOfType<KtDeclaration>(true) as? KtClassOrObject ?: return null
+    val leaf = when (val psi = location.psiElement) {
+        is KtLightElement<*, *> -> psi.kotlinOrigin
+        else -> psi
+    }
+    val owner = leaf?.getParentOfType<KtDeclaration>(false) as? KtClassOrObject ?: return null
     return KtFakeLightClass(owner)
 }
\ No newline at end of file