diff --git a/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/InvokeFunction.kt b/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/InvokeFunction.kt
index bc08f1a55fb2..eb184b21b745 100644
--- a/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/InvokeFunction.kt
+++ b/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/InvokeFunction.kt
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 import test1.MyClass
 import test1.invoke
 
diff --git a/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/InvokeFunction.kt.after b/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/InvokeFunction.kt.after
index bc08f1a55fb2..eb184b21b745 100644
--- a/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/InvokeFunction.kt.after
+++ b/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/InvokeFunction.kt.after
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 import test1.MyClass
 import test1.invoke
 
diff --git a/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/ResolvedImportAndUnresolvedReference.kt b/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/ResolvedImportAndUnresolvedReference.kt
index ba3cf1df6e67..1dd2b4fdfa02 100644
--- a/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/ResolvedImportAndUnresolvedReference.kt
+++ b/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/ResolvedImportAndUnresolvedReference.kt
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 // WITH_MESSAGE: "Unused imports not found"
 import one.doSmth
 
diff --git a/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/ResolvedImportAndUnresolvedReference.kt.after b/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/ResolvedImportAndUnresolvedReference.kt.after
index ba3cf1df6e67..1dd2b4fdfa02 100644
--- a/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/ResolvedImportAndUnresolvedReference.kt.after
+++ b/plugins/kotlin/idea/tests/testData/editor/optimizeImports/common/ResolvedImportAndUnresolvedReference.kt.after
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 // WITH_MESSAGE: "Unused imports not found"
 import one.doSmth