diff --git a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/callResolver/AbstractResolveCallTest.kt b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/callResolver/AbstractResolveCallTest.kt
index 5ea9a3e6788..4301262f706 100644
--- a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/callResolver/AbstractResolveCallTest.kt
+++ b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/callResolver/AbstractResolveCallTest.kt
@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.callResolver
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.calls.KtCallInfo
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.stringRepresentation
@@ -20,7 +19,7 @@ import org.jetbrains.kotlin.test.services.assertions
 abstract class AbstractResolveCallTest : AbstractAnalysisApiBasedSingleModuleTest() {
     override fun doTestByFileStructure(ktFiles: List<KtFile>, module: TestModule, testServices: TestServices) {
         val ktFile = ktFiles.first()
-        val expression = testServices.expressionMarkerProvider.getSelectedElement(ktFile)
+        val expression = testServices.expressionMarkerProvider.getSelectedElementOfType<KtElement>(ktFile)
 
         val actual = executeOnPooledThreadInReadAction {
             analyseForTest(expression) {
diff --git a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/callResolver/AbstractResolveCandidatesTest.kt b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/callResolver/AbstractResolveCandidatesTest.kt
index 200d71b755b..31318a7cdec 100644
--- a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/callResolver/AbstractResolveCandidatesTest.kt
+++ b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/callResolver/AbstractResolveCandidatesTest.kt
@@ -21,7 +21,7 @@ import org.jetbrains.kotlin.test.services.assertions
 abstract class AbstractResolveCandidatesTest : AbstractAnalysisApiBasedSingleModuleTest() {
     override fun doTestByFileStructure(ktFiles: List<KtFile>, module: TestModule, testServices: TestServices) {
         val ktFile = ktFiles.first()
-        val expression = testServices.expressionMarkerProvider.getSelectedElement(ktFile)
+        val expression = testServices.expressionMarkerProvider.getSelectedElementOfType<KtElement>(ktFile)
 
         val actual = executeOnPooledThreadInReadAction {
             analyseForTest(expression) {
diff --git a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/references/AbstractReferenceShortenerTest.kt b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/references/AbstractReferenceShortenerTest.kt
index 78db3a2bd49..15124acfe4f 100644
--- a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/references/AbstractReferenceShortenerTest.kt
+++ b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/references/AbstractReferenceShortenerTest.kt
@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.references.Shorten
 import org.jetbrains.kotlin.analysis.test.framework.base.AbstractAnalysisApiBasedSingleModuleTest
 import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
 import org.jetbrains.kotlin.analysis.test.framework.utils.executeOnPooledThreadInReadAction
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.model.TestModule
 import org.jetbrains.kotlin.test.services.TestServices
@@ -22,7 +23,7 @@ import org.jetbrains.kotlin.test.services.assertions
  */
 abstract class AbstractReferenceShortenerTest : AbstractAnalysisApiBasedSingleModuleTest() {
     override fun doTestByFileStructure(ktFiles: List<KtFile>, module: TestModule, testServices: TestServices) {
-        val element = testServices.expressionMarkerProvider.getSelectedElement(ktFiles.first())
+        val element = testServices.expressionMarkerProvider.getSelectedElementOfType<KtElement>(ktFiles.first())
 
         val shortenings = executeOnPooledThreadInReadAction {
             analyseForTest(element) {
diff --git a/analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/services/ExpressionMarkersSourceFilePreprocessor.kt b/analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/services/ExpressionMarkersSourceFilePreprocessor.kt
index efcd4648f50..2ae0fa433e8 100644
--- a/analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/services/ExpressionMarkersSourceFilePreprocessor.kt
+++ b/analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/services/ExpressionMarkersSourceFilePreprocessor.kt
@@ -124,14 +124,13 @@ class ExpressionMarkerProvider : TestService {
 
     }
 
-    fun getSelectedElement(file: KtFile): KtElement {
-        val range = selected[file.name]
-            ?: error("No selected expression found in file")
+    fun getSelectedElement(file: KtFile): PsiElement {
+        val range = selected[file.name] ?: error("No selected expression found in file")
         val elements = file.elementsInRange(range).trimWhitespaces()
         if (elements.size != 1) {
             error("Expected one element at rage but found ${elements.size} [${elements.joinToString { it::class.simpleName + ": " + it.text }}]")
         }
-        return elements.single() as KtElement
+        return elements.single()
     }
 
     fun getSelectedElementOfTypeByDirective(ktFile: KtFile, module: TestModule): PsiElement {
@@ -150,7 +149,7 @@ class ExpressionMarkerProvider : TestService {
     inline fun <reified E : KtElement> getSelectedElementOfType(file: KtFile): E {
         return when (val selected = getSelectedElement(file)) {
             is E -> selected
-            else -> generateSequence(selected as PsiElement) { current ->
+            else -> generateSequence(selected) { current ->
                 current.children.singleOrNull()?.takeIf { it.textRange == current.textRange }
             }.firstIsInstance()
         }