diff --git a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/scopeProvider/AbstractFileImportingScopeContextTest.kt b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/scopeProvider/AbstractFileImportingScopeContextTest.kt
index a1611670ec7..9bd693816a2 100644
--- a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/scopeProvider/AbstractFileImportingScopeContextTest.kt
+++ b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/scopeProvider/AbstractFileImportingScopeContextTest.kt
@@ -16,14 +16,13 @@ import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
 import org.jetbrains.kotlin.test.model.TestModule
 import org.jetbrains.kotlin.test.services.TestServices
 import org.jetbrains.kotlin.test.services.assertions
-import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 
 abstract class AbstractFileImportingScopeContextTest : AbstractAnalysisApiBasedSingleModuleTest() {
     override fun doTestByFileStructure(ktFiles: List<KtFile>, module: TestModule, testServices: TestServices) {
         val ktFile = ktFiles.first()
         val renderDefaultImportingScope = Directives.RENDER_DEFAULT_IMPORTING_SCOPE in module.directives
 
-        analyseForTest(ktFile.children.firstIsInstance()) {
+        analyseForTest(ktFile) {
             val ktScopeContext = ktFile.getImportingScopeContext()
 
             val scopeContextStringRepresentation = render(ktScopeContext, renderDefaultImportingScope)
diff --git a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/references/AbstractReferenceShortenerForWholeFileTest.kt b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/references/AbstractReferenceShortenerForWholeFileTest.kt
index 6a021424f9e..5bea526d2ef 100644
--- a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/references/AbstractReferenceShortenerForWholeFileTest.kt
+++ b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/references/AbstractReferenceShortenerForWholeFileTest.kt
@@ -18,7 +18,7 @@ abstract class AbstractReferenceShortenerForWholeFileTest : AbstractAnalysisApiB
         val file = ktFiles.first()
 
         val shortening = executeOnPooledThreadInReadAction {
-            analyseForTest(file.declarations.first()) {
+            analyseForTest(file) {
                 collectPossibleReferenceShortenings(file)
             }
         }
diff --git a/analysis/analysis-api/testData/components/referenceShortener/referenceShortenerWholeFile/fileWithMultipleDeclarations.kt b/analysis/analysis-api/testData/components/referenceShortener/referenceShortenerWholeFile/fileWithMultipleDeclarations.kt
index 3ab422ba922..798ead497ff 100644
--- a/analysis/analysis-api/testData/components/referenceShortener/referenceShortenerWholeFile/fileWithMultipleDeclarations.kt
+++ b/analysis/analysis-api/testData/components/referenceShortener/referenceShortenerWholeFile/fileWithMultipleDeclarations.kt
@@ -1,4 +1,4 @@
-class A {
+class A<caret_onAirContext> {
     abstract class B {
         abstract val b: A.B
     }
diff --git a/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/errorImport.kt b/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/errorImport.kt
index 1a598363abb..ac7ab9c22de 100644
--- a/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/errorImport.kt
+++ b/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/errorImport.kt
@@ -1,4 +1,4 @@
 import c.b
 
-val a = 5 // some KtElement to analyze in the air against
+val a<caret_onAirContext> = 5
 
diff --git a/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/importAlias.kt b/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/importAlias.kt
index 19a01c92ad9..6cc6f4b2007 100644
--- a/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/importAlias.kt
+++ b/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/importAlias.kt
@@ -1,7 +1,7 @@
 // FILE: main.kt
 import other.A as AAA
 
-val a = 5 // some KtElement to analyze in the air against
+val a<caret_onAirContext> = 5
 
 // FILE: other/other.kt
 package other
diff --git a/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/noImports.kt b/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/noImports.kt
index 925b1c7cf24..3d1977f2c98 100644
--- a/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/noImports.kt
+++ b/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/noImports.kt
@@ -1,3 +1,3 @@
 // RENDER_DEFAULT_IMPORTING_SCOPE
 
-val a = 5 // some KtElement to analyze in the air against
\ No newline at end of file
+val a<caret_onAirContext> = 5
diff --git a/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/simpleImport.kt b/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/simpleImport.kt
index ae829550c3b..98424560dd0 100644
--- a/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/simpleImport.kt
+++ b/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/simpleImport.kt
@@ -1,7 +1,7 @@
 // FILE: main.kt
 import other.A
 
-val a = 5 // some KtElement to analyze in the air against
+val a<caret_onAirContext> = 5
 
 // FILE: other/other.kt
 package other
diff --git a/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/starImport.kt b/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/starImport.kt
index ce2b18ceaf8..2d8c3e35268 100644
--- a/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/starImport.kt
+++ b/analysis/analysis-api/testData/components/scopeProvider/importingScopeContext/starImport.kt
@@ -1,7 +1,7 @@
 // FILE: main.kt
 import other.*
 
-val a = 5 // some KtElement to analyze in the air against
+val a<caret_onAirContext> = 5
 
 // FILE: other/other.kt
 package other
diff --git a/analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/base/AbstractAnalysisApiBasedTest.kt b/analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/base/AbstractAnalysisApiBasedTest.kt
index 5fe5e5dd4d9..616a5f4cb5e 100644
--- a/analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/base/AbstractAnalysisApiBasedTest.kt
+++ b/analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/base/AbstractAnalysisApiBasedTest.kt
@@ -17,11 +17,13 @@ import org.jetbrains.kotlin.analysis.test.framework.TestWithDisposable
 import org.jetbrains.kotlin.analysis.test.framework.project.structure.ktModuleProvider
 import org.jetbrains.kotlin.analysis.test.framework.services.ExpressionMarkerProvider
 import org.jetbrains.kotlin.analysis.test.framework.services.ExpressionMarkersSourceFilePreprocessor
+import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
 import org.jetbrains.kotlin.analysis.test.framework.services.libraries.CompilerExecutor
 import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestConfigurator
 import org.jetbrains.kotlin.analysis.test.framework.test.configurators.FrontendKind
 import org.jetbrains.kotlin.analysis.test.framework.utils.SkipTestException
 import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
+import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.TestConfiguration
@@ -192,6 +194,20 @@ abstract class AbstractAnalysisApiBasedTest : TestWithDisposable() {
         }
     }
 
+    /**
+     * Invoke the analysis in the context of given [file]
+     *
+     * To perform the test for in-air analysis, it will look for the declaration marked with the caret `<caret_onAirContext>`
+     */
+    protected fun <R> analyseForTest(file: KtFile, action: KtAnalysisSession.(KtElement) -> R): R {
+        return if (configurator.analyseInDependentSession) {
+            val declaration = testServices.expressionMarkerProvider.getElementOfTypeAtCaret<KtDeclaration>(file, ON_AIR_CONTEXT_CARET_TAG)
+            analyseForTest(declaration, action)
+        } else {
+            analyze(file, action = { action(file) })
+        }
+    }
+
     @BeforeEach
     fun initTestInfo(testInfo: TestInfo) {
         this.testInfo = KotlinTestInfo(
@@ -200,4 +216,8 @@ abstract class AbstractAnalysisApiBasedTest : TestWithDisposable() {
             tags = testInfo.tags
         )
     }
+
+    companion object {
+        private const val ON_AIR_CONTEXT_CARET_TAG = "onAirContext"
+    }
 }
\ No newline at end of file