diff --git a/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/AbstractDiagnosticsTest.kt b/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/AbstractDiagnosticsTest.kt
index a2e1c5bfb47..c806bbb8d78 100644
--- a/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/AbstractDiagnosticsTest.kt
+++ b/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/AbstractDiagnosticsTest.kt
@@ -142,12 +142,13 @@ abstract class AbstractDiagnosticsTest : BaseDiagnosticsTest() {
         var exceptionFromDescriptorValidation: Throwable? = null
         val originalTestFile = testDataFile.readText()
         try {
-            val isJavacExpectedFile = InTextDirectivesUtils.isDirectiveDefined(originalTestFile, "// JAVAC_EXPECTED_FILE")
-                                      && environment.configuration.getBoolean(JVMConfigurationKeys.USE_JAVAC)
-
             val postfix = when {
-                isJavacExpectedFile -> ".javac.txt"
-                files.all { it.withNewInferenceDirective && it.newInferenceEnabled } -> ".ni.txt"
+                InTextDirectivesUtils.isDirectiveDefined(originalTestFile, "// JAVAC_EXPECTED_FILE") &&
+                environment.configuration.getBoolean(JVMConfigurationKeys.USE_JAVAC) -> ".javac.txt"
+
+                InTextDirectivesUtils.isDirectiveDefined(originalTestFile, "// NI_EXPECTED_FILE") &&
+                files.any { it.newInferenceEnabled } && !USE_OLD_INFERENCE_DIAGNOSTICS_FOR_NI -> ".ni.txt"
+
                 else -> ".txt"
             }
             val expectedFile = File(FileUtil.getNameWithoutExtension(testDataFile.absolutePath) + postfix)
diff --git a/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/BaseDiagnosticsTest.kt b/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/BaseDiagnosticsTest.kt
index 2cbab157ad6..09e38b274e4 100644
--- a/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/BaseDiagnosticsTest.kt
+++ b/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/BaseDiagnosticsTest.kt
@@ -325,7 +325,7 @@ abstract class BaseDiagnosticsTest : KotlinMultiFileTestWithJava<TestModule, Tes
         val WITH_NEW_INFERENCE_DIRECTIVE = "WITH_NEW_INFERENCE"
 
         // Change it to "true" to load diagnostics for old inference to test new inference (ignore diagnostics with <NI; prefix)
-        private val USE_OLD_INFERENCE_DIAGNOSTICS_FOR_NI = false
+        val USE_OLD_INFERENCE_DIAGNOSTICS_FOR_NI = false
 
         private fun parseDiagnosticFilterDirective(directiveMap: Map<String, String>, allowUnderscoreUsage: Boolean): Condition<Diagnostic> {
             val directives = directiveMap[DIAGNOSTICS_DIRECTIVE]