diff --git a/compiler/fir/raw-fir/light-tree2fir/tests/org/jetbrains/kotlin/fir/lightTree/compare/TreesCompareTest.kt b/compiler/fir/raw-fir/light-tree2fir/tests/org/jetbrains/kotlin/fir/lightTree/compare/TreesCompareTest.kt
index 2f983a0f9b5..4d43c8dd4c2 100644
--- a/compiler/fir/raw-fir/light-tree2fir/tests/org/jetbrains/kotlin/fir/lightTree/compare/TreesCompareTest.kt
+++ b/compiler/fir/raw-fir/light-tree2fir/tests/org/jetbrains/kotlin/fir/lightTree/compare/TreesCompareTest.kt
@@ -68,19 +68,25 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
             val (text, linesMapping) = with(file.inputStream().reader(Charsets.UTF_8)) {
                 this.readSourceFileWithMapping()
             }
+            splitText(file.path, text.toString().trim()).forEach { pair ->
+                val (filePath, fileText) = pair
 
-            //psi
-            val ktFile = createPsiFile(FileUtil.getNameWithoutExtension(PathUtil.getFileName(file.path)), text.toString().trim()) as KtFile
-            val firFileFromPsi = ktFile.toFirFile()
-            val treeFromPsi = FirRenderer().renderElementAsString(firFileFromPsi)
-                .replace("<ERROR TYPE REF:.*?>".toRegex(), "<ERROR TYPE REF>")
+                //psi
+                val ktFile = createPsiFile(FileUtil.getNameWithoutExtension(PathUtil.getFileName(filePath)), fileText) as KtFile
+                val firFileFromPsi = ktFile.toFirFile()
+                val treeFromPsi = FirRenderer().renderElementAsString(firFileFromPsi)
+                    .replace("<ERROR TYPE REF:.*?>".toRegex(), "<ERROR TYPE REF>")
 
-            //light tree
-            val firFileFromLightTree = lightTreeConverter.buildFirFile(text, KtIoFileSourceFile(file), linesMapping)
-            val treeFromLightTree = FirRenderer().renderElementAsString(firFileFromLightTree)
-                .replace("<ERROR TYPE REF:.*?>".toRegex(), "<ERROR TYPE REF>")
+                //light tree
+                val firFileFromLightTree = lightTreeConverter.buildFirFile(text, KtIoFileSourceFile(file), linesMapping)
+                val treeFromLightTree = FirRenderer().renderElementAsString(firFileFromLightTree)
+                    .replace("<ERROR TYPE REF:.*?>".toRegex(), "<ERROR TYPE REF>")
 
-            return@compareBase treeFromLightTree == treeFromPsi
+                if (treeFromLightTree != treeFromPsi) {
+                    return@compareBase false
+                }
+            }
+            return@compareBase true
         }
     }
 
@@ -101,29 +107,56 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
             val notEditedText = FileUtil.loadFile(file, CharsetToolkit.UTF8, true).trim()
             val text = notEditedText.replace(DIAGNOSTIC_IN_TESTDATA_PATTERN, "").replaceAfter(".java", "")
 
-            //psi
-            val ktFile = createPsiFile(FileUtil.getNameWithoutExtension(PathUtil.getFileName(file.path)), text) as KtFile
-            val firFileFromPsi = ktFile.toFirFile()
-            val treeFromPsi = FirRenderer().renderElementAsString(firFileFromPsi)
-                .replace("<Unsupported LValue.*?>".toRegex(), "<Unsupported LValue>")
-                .replace("<ERROR TYPE REF:.*?>".toRegex(), "<ERROR TYPE REF>")
+            splitText(file.path, text).forEach { pair ->
+                val (filePath, fileText) = pair
+                //psi
+                val fileName = PathUtil.getFileName(filePath)
+                val ktFile = createPsiFile(FileUtil.getNameWithoutExtension(fileName), fileText) as KtFile
+                val firFileFromPsi = ktFile.toFirFile()
+                val treeFromPsi = FirRenderer().renderElementAsString(firFileFromPsi)
+                    .replace("<Unsupported LValue.*?>".toRegex(), "<Unsupported LValue>")
+                    .replace("<ERROR TYPE REF:.*?>".toRegex(), "<ERROR TYPE REF>")
 
-            //light tree
-            val firFileFromLightTree =
-                lightTreeConverter.buildFirFile(
-                    text,
-                    KtInMemoryTextSourceFile(file.name, file.path, text),
-                    text.toSourceLinesMapping()
-                )
-            val treeFromLightTree = FirRenderer().renderElementAsString(firFileFromLightTree)
-                .replace("<Unsupported LValue.*?>".toRegex(), "<Unsupported LValue>")
-                .replace("<ERROR TYPE REF:.*?>".toRegex(), "<ERROR TYPE REF>")
+                //light tree
+                val firFileFromLightTree =
+                    lightTreeConverter.buildFirFile(
+                        fileText,
+                        KtInMemoryTextSourceFile(fileName, filePath, fileText),
+                        fileText.toSourceLinesMapping()
+                    )
+                val treeFromLightTree = FirRenderer().renderElementAsString(firFileFromLightTree)
+                    .replace("<Unsupported LValue.*?>".toRegex(), "<Unsupported LValue>")
+                    .replace("<ERROR TYPE REF:.*?>".toRegex(), "<ERROR TYPE REF>")
 
-            return@compareBase treeFromLightTree == treeFromPsi
+                if (treeFromLightTree != treeFromPsi) {
+                    return@compareBase false
+                }
+            }
+            return@compareBase true
         }
     }
 
     fun testCompareAll() {
         compareAll()
     }
+
+    private fun splitText(filePath: String, text: String): List<Pair<String, String>> {
+        val fileDirective = "// FILE:"
+        val idx = text.indexOf(fileDirective)
+        if (idx > 0 && text[idx - 1] != '\n') {
+            //try to avoid splitting of sources
+            return emptyList()
+        }
+        if (idx >= 0) {
+            val result = mutableListOf<Pair<String, String>>()
+            val strings = text.drop(idx).drop(fileDirective.length).split(fileDirective)
+            for (string in strings) {
+                val newLineIdx = string.indexOf("\n")
+                if (newLineIdx < 0) return emptyList()
+                result.add(Pair(string.substring(0, newLineIdx).trim(), string.substring(newLineIdx)))
+            }
+            return result
+        }
+        return listOf(Pair(filePath, text))
+    }
 }