From 12cd5fb43af76f7cc09c4d03136fb1e1b0cb16a0 Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Thu, 1 Oct 2020 13:02:51 +0300
Subject: [PATCH] [FIR Completion] Refactor FIR completion tests

Extract logic about checking the directive and catching the exception
---
 ...AbstractHighLevelJvmBasicCompletionTest.kt | 11 +----
 .../highLevelCompletionTestUtils.kt           | 44 +++++++++++++++++++
 2 files changed, 45 insertions(+), 10 deletions(-)
 create mode 100644 idea/idea-fir/tests/org/jetbrains/kotlin/idea/completion/highLevelCompletionTestUtils.kt

diff --git a/idea/idea-fir/tests/org/jetbrains/kotlin/idea/completion/AbstractHighLevelJvmBasicCompletionTest.kt b/idea/idea-fir/tests/org/jetbrains/kotlin/idea/completion/AbstractHighLevelJvmBasicCompletionTest.kt
index 3930cb796e7..3e732678647 100644
--- a/idea/idea-fir/tests/org/jetbrains/kotlin/idea/completion/AbstractHighLevelJvmBasicCompletionTest.kt
+++ b/idea/idea-fir/tests/org/jetbrains/kotlin/idea/completion/AbstractHighLevelJvmBasicCompletionTest.kt
@@ -12,15 +12,6 @@ abstract class AbstractHighLevelJvmBasicCompletionTest : AbstractJvmBasicComplet
     override val captureExceptions: Boolean = false
 
     override fun executeTest(test: () -> Unit) {
-        val doComparison = InTextDirectivesUtils.isDirectiveDefined(myFixture.file.text, "FIR_COMPARISON")
-        try {
-            test()
-        } catch (e: Throwable) {
-            if (doComparison) throw e
-            return
-        }
-        if (!doComparison) {
-            throw AssertionError("Looks like test is passing, please add // FIR_COMPARISON at the beginning of the file")
-        }
+        runTestWithCustomEnableDirective(FIR_COMPARISON, testDataFile()) { super.executeTest(test) }
     }
 }
\ No newline at end of file
diff --git a/idea/idea-fir/tests/org/jetbrains/kotlin/idea/completion/highLevelCompletionTestUtils.kt b/idea/idea-fir/tests/org/jetbrains/kotlin/idea/completion/highLevelCompletionTestUtils.kt
new file mode 100644
index 00000000000..26edc6a77ef
--- /dev/null
+++ b/idea/idea-fir/tests/org/jetbrains/kotlin/idea/completion/highLevelCompletionTestUtils.kt
@@ -0,0 +1,44 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.completion
+
+import org.jetbrains.kotlin.test.InTextDirectivesUtils
+import java.io.File
+
+const val FIR_COMPARISON = "// FIR_COMPARISON"
+
+/**
+ * Set this flag to `true` to insert directive automatically to all files
+ * that pass tests but do not already have the directive.
+ */
+private const val insertDirectiveAutomatically = false
+
+fun runTestWithCustomEnableDirective(directive: String, testFile: File, test: () -> Unit) {
+    val testFileAfter = testFile.resolveSibling(testFile.name + ".after").takeIf { it.exists() }
+
+    val testEnabled = InTextDirectivesUtils.isDirectiveDefined(testFile.readText(), directive)
+
+    try {
+        test()
+    } catch (e: Throwable) {
+        if (testEnabled) throw e
+        return
+    }
+
+    if (!testEnabled) {
+        if (insertDirectiveAutomatically) {
+            testFile.insertDirective(directive)
+            testFileAfter?.insertDirective(directive)
+        }
+
+        throw AssertionError("Looks like test is passing, please add ${directive.removePrefix("// ")} at the beginning of the file")
+    }
+}
+
+private fun File.insertDirective(directive: String) {
+    val originalText = readText()
+    writeText("$directive\n$originalText")
+}