From d5d18b2ad67e497db99fa5c8b7c471884dfb1c36 Mon Sep 17 00:00:00 2001
From: Denis Zharkov <denis.zharkov@jetbrains.com>
Date: Wed, 28 Sep 2016 17:58:58 +0300
Subject: [PATCH] Enhance tests for postfix templates

Prohibit multiple suggested expressions in test mode
---
 .../postfix/KtPostfixTemplateProvider.kt      | 21 +++++++++++++++++++
 .../AbstractPostfixTemplateProviderTest.kt    | 11 ++++++++++
 2 files changed, 32 insertions(+)

diff --git a/idea/src/org/jetbrains/kotlin/idea/codeInsight/postfix/KtPostfixTemplateProvider.kt b/idea/src/org/jetbrains/kotlin/idea/codeInsight/postfix/KtPostfixTemplateProvider.kt
index bb92a3f6c05..39b4697eeac 100644
--- a/idea/src/org/jetbrains/kotlin/idea/codeInsight/postfix/KtPostfixTemplateProvider.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/codeInsight/postfix/KtPostfixTemplateProvider.kt
@@ -17,11 +17,13 @@
 package org.jetbrains.kotlin.idea.codeInsight.postfix
 
 import com.intellij.codeInsight.template.postfix.templates.*
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.Condition
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.intentions.negate
 import org.jetbrains.kotlin.idea.refactoring.introduce.introduceVariable.KotlinIntroduceVariableHandler
@@ -63,6 +65,15 @@ class KtPostfixTemplateProvider : PostfixTemplateProvider {
 
     override fun preExpand(file: PsiFile, editor: Editor) {
     }
+
+    companion object {
+        /**
+         * In tests only one expression should be suggested, so in case there are many of them, save relevant items
+         */
+        @TestOnly
+        @Volatile
+        var previouslySuggestedExpressions = emptyList<String>()
+    }
 }
 
 private object KtNotPostfixTemplate : NotPostfixTemplate(
@@ -108,6 +119,16 @@ private fun PsiElement.isStatement() = parent is KtBlockExpression
 private class KtExpressionPostfixTemplateSelector(
         filter: PsiElement.() -> Boolean
 ) : PostfixTemplateExpressionSelectorBase(Condition(filter)) {
+    override fun getExpressions(context: PsiElement, document: Document, offset: Int): List<PsiElement> {
+        val expressions = super.getExpressions(context, document, offset)
+
+        if (ApplicationManager.getApplication().isUnitTestMode && expressions.size > 1) {
+            KtPostfixTemplateProvider.previouslySuggestedExpressions = expressions.map { it.text }
+        }
+
+        return expressions
+    }
+
     override fun getNonFilteredExpressions(
             context: PsiElement,
             document: Document,
diff --git a/idea/tests/org/jetbrains/kotlin/idea/codeInsight/postfix/AbstractPostfixTemplateProviderTest.kt b/idea/tests/org/jetbrains/kotlin/idea/codeInsight/postfix/AbstractPostfixTemplateProviderTest.kt
index 826ba4e6f5d..918a29fa45b 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/codeInsight/postfix/AbstractPostfixTemplateProviderTest.kt
+++ b/idea/tests/org/jetbrains/kotlin/idea/codeInsight/postfix/AbstractPostfixTemplateProviderTest.kt
@@ -29,6 +29,17 @@ abstract class AbstractPostfixTemplateProviderTest : KotlinLightCodeInsightFixtu
     protected fun doTest(fileName: String) {
         myFixture.configureByFile(fileName)
         myFixture.type("\t")
+
+        val previouslySuggestedExpressions = KtPostfixTemplateProvider.previouslySuggestedExpressions
+        if (previouslySuggestedExpressions.size > 1) {
+            fail("Only one expression should be suggested, but $previouslySuggestedExpressions were found")
+        }
+
         myFixture.checkResultByFile(fileName + ".after")
     }
+
+    override fun tearDown() {
+        super.tearDown()
+        KtPostfixTemplateProvider.previouslySuggestedExpressions = emptyList()
+    }
 }