diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirOnAirResolveTest.kt b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirOnAirResolveTest.kt
index d8a8f144bfc..3de33bd1deb 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirOnAirResolveTest.kt
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirOnAirResolveTest.kt
@@ -6,9 +6,9 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LowLevelFirApiFacadeForResolveOnAir
-import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirSourceResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirSourceTestConfigurator
+import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.fir.renderer.FirRenderer
 import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.KtElement
@@ -35,7 +35,7 @@ abstract class AbstractFirOnAirResolveTest : AbstractLowLevelApiSingleFileTest()
         check(place::class == onAir::class)
 
         resolveWithClearCaches(ktFile) { firResolveSession ->
-            check(firResolveSession is LLFirSourceResolveSession)
+            check(firResolveSession.useSiteKtModule is KtSourceModule)
             val firElement = LowLevelFirApiFacadeForResolveOnAir.onAirResolveElement(firResolveSession, place, onAir)
             val rendered = FirRenderer.withResolvePhase().renderElementAsString(firElement)
             testServices.assertions.assertEqualsToTestDataFileSibling(rendered)
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/firTestUtils.kt b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/firTestUtils.kt
index fb3b248cc80..f7d37cad704 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/firTestUtils.kt
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/firTestUtils.kt
@@ -7,9 +7,8 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionConfigurator
-import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirNotUnderContentRootResolveSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirScriptResolveSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirSourceResolveSession
+import org.jetbrains.kotlin.analysis.project.structure.KtLibraryModule
+import org.jetbrains.kotlin.analysis.project.structure.KtLibrarySourceModule
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.test.framework.services.environmentManager
 import org.jetbrains.kotlin.fir.FirElement
@@ -54,11 +53,9 @@ internal inline fun <R> resolveWithCaches(context: KtElement, action: (LLFirReso
 }
 
 internal val LLFirResolveSession.isSourceSession: Boolean
-    get() {
-        return when (this) {
-            is LLFirSourceResolveSession, is LLFirNotUnderContentRootResolveSession, is LLFirScriptResolveSession -> true
-            else -> false
-        }
+    get() = when (useSiteKtModule) {
+        is KtLibraryModule, is KtLibrarySourceModule -> false
+        else -> true
     }
 
 internal fun TestConfigurationBuilder.useFirSessionConfigurator(configurator: (TestServices) -> LLFirSessionConfigurator) {