[Test] Correcly collect diagnostics in AA tests

One of the previous commits moved calculation of diagnostics from
  FirAnalyzerFacadeinto the separate test service, so it should be altered
  in AA tests too

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/compiler/based/AbstractCompilerBasedTestForFir.kt

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.test.frontend.fir.FirFrontendFacade
 import org.jetbrains.kotlin.test.frontend.fir.FirOutputArtifact
 import org.jetbrains.kotlin.test.frontend.fir.FirOutputArtifactImpl
 import org.jetbrains.kotlin.test.frontend.fir.FirOutputPartForDependsOnModule
+import org.jetbrains.kotlin.test.frontend.fir.handlers.FirDiagnosticCollectorService
 import org.jetbrains.kotlin.test.model.DependencyKind
 import org.jetbrains.kotlin.test.model.FrontendKinds
 import org.jetbrains.kotlin.test.model.TestModule
@@ -52,6 +53,7 @@ abstract class AbstractCompilerBasedTestForFir : AbstractCompilerBasedTest() {
         registerAnalysisApiBaseTestServices(disposable, FirLowLevelCompilerBasedTestConfigurator)
         useDirectives(SealedClassesInheritorsCaclulatorPreAnalysisHandler.Directives)
         usePreAnalysisHandlers(::SealedClassesInheritorsCaclulatorPreAnalysisHandler)
+        useAdditionalServices(service<FirDiagnosticCollectorService>(::AnalysisApiFirDiagnosticCollectorService))
 
         firHandlersStep {
             useHandlers(::LLDiagnosticParameterChecker)
@@ -127,4 +129,4 @@ abstract class AbstractCompilerBasedTestForFir : AbstractCompilerBasedTest() {
 
         super.runTest(filePath)
     }
-}
+}

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/compiler/based/LLDiagnosticParameterChecker.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.compiler.based
 
-import org.jetbrains.kotlin.analysis.utils.errors.withKtModuleEntry
 import org.jetbrains.kotlin.diagnostics.*
 import org.jetbrains.kotlin.diagnostics.rendering.RootDiagnosticRendererFactory
 import org.jetbrains.kotlin.fir.types.ConeKotlinType

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/compiler/based/LowLevelFirAnalyzerFacade.kt

@@ -16,8 +16,14 @@ import org.jetbrains.kotlin.fir.pipeline.FirResult
 import org.jetbrains.kotlin.fir.pipeline.ModuleCompilerAnalyzedOutput
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.util.listMultimapOf
+import org.jetbrains.kotlin.fir.util.plusAssign
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.test.frontend.fir.FirOutputArtifact
+import org.jetbrains.kotlin.test.frontend.fir.handlers.DiagnosticsMap
+import org.jetbrains.kotlin.test.frontend.fir.handlers.FirDiagnosticCollectorService
 import org.jetbrains.kotlin.test.model.TestFile
+import org.jetbrains.kotlin.test.services.TestServices
 
 open class LowLevelFirAnalyzerFacade(
     val firResolveSession: LLFirResolveSession,
@@ -35,7 +41,6 @@ open class LowLevelFirAnalyzerFacade(
 
     private var resolved: Boolean = false
 
-    // TODO: investigate
     fun runCheckers(): Map<FirFile, List<KtDiagnostic>> {
         if (!resolved) {
             runResolution()
@@ -52,3 +57,15 @@ open class LowLevelFirAnalyzerFacade(
 
     override fun runResolution(): List<FirFile> = allFirFiles.values.toList()
 }
+
+class AnalysisApiFirDiagnosticCollectorService(testServices: TestServices) : FirDiagnosticCollectorService(testServices) {
+    override fun getFrontendDiagnosticsForModule(info: FirOutputArtifact): DiagnosticsMap {
+        val result = listMultimapOf<FirFile, KtDiagnostic>()
+        for (part in info.partsForDependsOnModules) {
+            val facade = part.firAnalyzerFacade
+            require(facade is LowLevelFirAnalyzerFacade)
+            result += facade.runCheckers()
+        }
+        return result
+    }
+}

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/handlers/FirDiagnosticsHandler.kt

@@ -61,7 +61,6 @@ import org.jetbrains.kotlin.test.utils.AbstractTwoAttributesMetaInfoProcessor
 import org.jetbrains.kotlin.test.utils.MultiModuleInfoDumper
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.util.capitalizeDecapitalize.toLowerCaseAsciiOnly
-import org.jetbrains.kotlin.utils.addToStdlib.getOrPut
 
 class FullDiagnosticsRenderer(private val directive: SimpleDirective) {
     private val dumper: MultiModuleInfoDumper = MultiModuleInfoDumper(moduleHeaderTemplate = "// -- Module: <%s> --")
@@ -605,10 +604,10 @@ fun KtDiagnostic.toMetaInfos(
 
 typealias DiagnosticsMap = Multimap<FirFile, KtDiagnostic, List<KtDiagnostic>>
 
-class FirDiagnosticCollectorService(val testServices: TestServices) : TestService {
+open class FirDiagnosticCollectorService(val testServices: TestServices) : TestService {
     private val cache: MutableMap<FirOutputArtifact, DiagnosticsMap> = mutableMapOf()
 
-    fun getFrontendDiagnosticsForModule(info: FirOutputArtifact): DiagnosticsMap {
+    open fun getFrontendDiagnosticsForModule(info: FirOutputArtifact): DiagnosticsMap {
         return cache.getOrPut(info) { computeDiagnostics(info) }
     }