[FIR] Add ability to register extension points in diagnostic tests

compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirBaseDiagnosticsTest.kt

@@ -27,6 +27,8 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.extensions.FirExtensionPointService
+import org.jetbrains.kotlin.fir.extensions.extensionPointService
 import org.jetbrains.kotlin.fir.java.FirJavaModuleBasedSession
 import org.jetbrains.kotlin.fir.java.FirLibrarySession
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
@@ -76,14 +78,18 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
         FirLibrarySession.create(
             builtInsModuleInfo, sessionProvider, allProjectScope, project,
             environment.createPackagePartProvider(allProjectScope)
-        )
+        ).also {
+            registerFirExtensions(it.extensionPointService)
+        }
 
         val configToSession = modules.mapValues { (config, info) ->
             val moduleFiles = groupedByModule.getValue(config)
             val scope = TopDownAnalyzerFacadeForJVM.newModuleSearchScope(
                 project,
                 moduleFiles.mapNotNull { it.ktFile })
-            FirJavaModuleBasedSession(info, sessionProvider, scope)
+            FirJavaModuleBasedSession(info, sessionProvider, scope).also {
+                registerFirExtensions(it.extensionPointService)
+            }
         }
 
         val firFilesPerSession = mutableMapOf<FirSession, List<FirFile>>()
@@ -102,6 +108,8 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
         runAnalysis(testDataFile, files, firFilesPerSession)
     }
 
+    open fun registerFirExtensions(service: FirExtensionPointService) {}
+
     private fun mapKtFilesToFirFiles(session: FirSession, ktFiles: List<KtFile>, firFiles: MutableList<FirFile>, useLightTree: Boolean) {
         val firProvider = (session.firProvider as FirProviderImpl)
         if (useLightTree) {