[Analysis API] tests: add possibility to specify LOOK_UP_FOR_ELEMENT_OF_TYPE for all tests

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/base/AbstractAnalysisApiBasedTest.kt

@@ -114,6 +114,7 @@ abstract class AbstractAnalysisApiBasedTest : TestWithDisposable() {
 
         useSourcePreprocessor(::ExpressionMarkersSourceFilePreprocessor)
         useAdditionalService { ExpressionMarkerProvider() }
+        useDirectives(ExpressionMarkerProvider.Directives)
 
         registerAnalysisApiBaseTestServices(disposable, configurator)
         configureTest(this)

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/services/ExpressionMarkersSourceFilePreprocessor.kt

@@ -13,8 +13,11 @@ import org.jetbrains.kotlin.analysis.utils.printer.parentOfType
 import org.jetbrains.kotlin.fir.PrivateForInline
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
 import org.jetbrains.kotlin.psi.psiUtil.elementsInRange
+import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
 import org.jetbrains.kotlin.test.model.TestFile
+import org.jetbrains.kotlin.test.model.TestModule
 import org.jetbrains.kotlin.test.services.SourceFilePreprocessor
 import org.jetbrains.kotlin.test.services.TestModuleStructure
 import org.jetbrains.kotlin.test.services.TestService
@@ -131,6 +134,17 @@ class ExpressionMarkerProvider : TestService {
         return elements.single() as KtElement
     }
 
+    fun getSelectedElementOfTypeByDirective(ktFile: KtFile, module: TestModule): PsiElement {
+        val selectedElement = getSelectedElement(ktFile)
+        val expectedType = module.directives[Directives.LOOK_UP_FOR_ELEMENT_OF_TYPE].firstOrNull() ?: return selectedElement
+        @Suppress("UNCHECKED_CAST") val expectedClass = Class.forName(expectedType) as Class<PsiElement>
+        if (expectedClass.isInstance(selectedElement)) return selectedElement
+
+        return selectedElement.collectDescendantsOfType<PsiElement> {
+            expectedClass.isInstance(it)
+        }.single { it.textRange == selectedElement.textRange }
+    }
+
     inline fun <reified E : KtElement> getSelectedElementOfType(file: KtFile): E {
         return when (val selected = getSelectedElement(file)) {
             is E -> selected
@@ -143,6 +157,10 @@ class ExpressionMarkerProvider : TestService {
     private fun List<PsiElement>.trimWhitespaces(): List<PsiElement> =
         dropWhile { it is PsiWhiteSpace }
             .dropLastWhile { it is PsiWhiteSpace }
+
+    object Directives : SimpleDirectivesContainer() {
+        val LOOK_UP_FOR_ELEMENT_OF_TYPE by stringDirective("LOOK_UP_FOR_ELEMENT_OF_TYPE")
+    }
 }
 
 @PrivateForInline

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractGetOrBuildFirTest.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirOutOfContentRootTestConfigurator
@@ -19,23 +18,15 @@ import org.jetbrains.kotlin.fir.renderer.FirRenderer
 import org.jetbrains.kotlin.fir.renderer.FirResolvePhaseRenderer
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
-import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
-import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
 import org.jetbrains.kotlin.test.services.TestModuleStructure
 import org.jetbrains.kotlin.test.services.TestServices
 import org.jetbrains.kotlin.test.services.assertions
 
 abstract class AbstractGetOrBuildFirTest : AbstractLowLevelApiSingleFileTest() {
-    override fun configureTest(builder: TestConfigurationBuilder) {
-        super.configureTest(builder)
-        with(builder) {
-            useDirectives(Directives)
-        }
-    }
 
     override fun doTestByFileStructure(ktFile: KtFile, moduleStructure: TestModuleStructure, testServices: TestServices) {
-        val selectedElement = getElementOfType(ktFile, moduleStructure, testServices) as KtElement
+        val module = moduleStructure.modules.single()
+        val selectedElement = testServices.expressionMarkerProvider.getSelectedElementOfTypeByDirective(ktFile, module) as KtElement
 
         val actual = resolveWithClearCaches(ktFile) { state ->
             val fir = selectedElement.getOrBuildFir(state)
@@ -49,17 +40,6 @@ abstract class AbstractGetOrBuildFirTest : AbstractLowLevelApiSingleFileTest() {
         testServices.assertions.assertEqualsToTestDataFileSibling(actual)
     }
 
-    private fun getElementOfType(ktFile: KtFile, moduleStructure: TestModuleStructure, testServices: TestServices): PsiElement {
-        val selectedElement = testServices.expressionMarkerProvider.getSelectedElement(ktFile)
-        val expectedType = moduleStructure.allDirectives[Directives.LOOK_UP_FOR_ELEMENT_OF_TYPE].firstOrNull() ?: return selectedElement
-        @Suppress("UNCHECKED_CAST") val expectedClass = Class.forName(expectedType) as Class<PsiElement>
-        if (expectedClass.isInstance(selectedElement)) return selectedElement
-
-        return selectedElement.collectDescendantsOfType<PsiElement> {
-            expectedClass.isInstance(it)
-        }.single { it.textRange == selectedElement.textRange }
-    }
-
     private fun render(firElement: FirElement?): String = when (firElement) {
         null -> "null"
         is FirImport -> "import ${firElement.importedFqName}"
@@ -69,10 +49,6 @@ abstract class AbstractGetOrBuildFirTest : AbstractLowLevelApiSingleFileTest() {
             resolvePhaseRenderer = FirResolvePhaseRenderer(),
         ).renderElementAsString(firElement)
     }
-
-    private object Directives : SimpleDirectivesContainer() {
-        val LOOK_UP_FOR_ELEMENT_OF_TYPE by stringDirective("LOOK_UP_FOR_ELEMENT_OF_TYPE")
-    }
 }
 
 abstract class AbstractSourceGetOrBuildFirTest : AbstractGetOrBuildFirTest() {