[FIR IDE] Fixed invalid leaks test

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/FirRefWithValidityCheck.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir.utils
 
+import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
@@ -19,6 +20,10 @@ internal class FirRefWithValidityCheck<out D : FirDeclaration>(fir: D, resolveSt
     private val firWeakRef = WeakReference(fir)
     private val resolveStateWeakRef = WeakReference(resolveState)
 
+    @TestOnly
+    internal fun isCollected(): Boolean =
+        firWeakRef.get() == null && resolveStateWeakRef.get() == null
+
     inline fun <R> withFir(phase: FirResolvePhase = FirResolvePhase.RAW_FIR, crossinline action: (fir: D) -> R): R {
         token.assertIsValid()
         val fir = firWeakRef.get()

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/frontend/api/symbols/AbstractMemoryLeakInSymbolsTest.kt

@@ -58,11 +58,8 @@ abstract class AbstractMemoryLeakInSymbolsTest : KotlinLightCodeInsightFixtureTe
 
     private fun KtSymbol.hasNoFirElementLeak(): LeakCheckResult {
         require(this is KtFirSymbol<*>)
-        return try {
+        return if (firRef.isCollected()) LeakCheckResult.NoLeak
-            firRef.withFir { LeakCheckResult.Leak(this::class.simpleName!!) }
+        else LeakCheckResult.Leak(this::class.simpleName!!)
-        } catch (_: EntityWasGarbageCollectedException) {
-            LeakCheckResult.NoLeak
-        }
     }
 
     private sealed class LeakCheckResult {