Use dialog wrapper for controlling behaviour from tests

Known issue: can't use same approach for modifying Java dialogs
behaviour

idea/idea-core/src/org/jetbrains/kotlin/idea/core/util/messages.kt

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.core.util
+
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.ui.Messages
+import org.jetbrains.annotations.TestOnly
+import java.util.concurrent.ConcurrentHashMap
+import javax.swing.Icon
+
+fun showYesNoCancelDialog(key: String, project: Project, message: String, title: String, icon: Icon, default: Int?): Int {
+    return if (!ApplicationManager.getApplication().isUnitTestMode) {
+        Messages.showYesNoCancelDialog(project, message, message, icon)
+    }
+    else {
+        callInTestMode(key, default)
+    }
+}
+
+private val dialogResults = ConcurrentHashMap<String, Any>()
+
+@TestOnly
+fun setDialogsResult(key: String, result: Any) {
+    dialogResults[key] = result
+}
+
+@TestOnly
+fun clearDialogsResults() {
+    dialogResults.clear()
+}
+
+private fun <T: Any?> callInTestMode(key: String, default: T?): T {
+    val result = dialogResults[key]
+    if (result != null) {
+        dialogResults.remove(key)
+        @Suppress("UNCHECKED_CAST")
+        return result as T
+    }
+
+    if (default != null) {
+        return default
+    }
+
+    throw IllegalStateException("Can't call '$key' dialog in test mode")
+}
+

idea/src/org/jetbrains/kotlin/idea/refactoring/kotlinRefactoringUtil.kt

@@ -74,6 +74,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.getJavaMemberDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
 import org.jetbrains.kotlin.idea.core.*
+import org.jetbrains.kotlin.idea.core.util.showYesNoCancelDialog
 import org.jetbrains.kotlin.idea.highlighter.markers.getAccessorLightMethods
 import org.jetbrains.kotlin.idea.intentions.RemoveCurlyBracesFromTemplateIntention
 import org.jetbrains.kotlin.idea.j2k.IdeaJavaToKotlinServices
@@ -105,6 +106,8 @@ import java.lang.annotation.Retention
 import java.util.*
 import javax.swing.Icon
 
+val CHECK_SUPER_METHODS_YES_NO_DIALOG = "CHECK_SUPER_METHODS_YES_NO_DIALOG"
+
 @JvmOverloads
 fun getOrCreateKotlinFile(fileName: String,
                           targetDir: PsiDirectory,
@@ -854,8 +857,6 @@ fun checkSuperMethods(
             declarationDescriptor: CallableDescriptor,
             overriddenElementsToDescriptor: Map<PsiElement, CallableDescriptor>
     ): List<PsiElement> {
-        if (ApplicationManager.getApplication().isUnitTestMode) return overriddenElementsToDescriptor.keys.toList()
-
         val superClassDescriptions = getClassDescriptions(overriddenElementsToDescriptor)
 
         val message = KotlinBundle.message(
@@ -865,7 +866,9 @@ fun checkSuperMethods(
                 actionString
         )
 
-        val exitCode = Messages.showYesNoCancelDialog(declaration.project, message, IdeBundle.message("title.warning"), Messages.getQuestionIcon())
+        val exitCode = showYesNoCancelDialog(
+                CHECK_SUPER_METHODS_YES_NO_DIALOG,
+                declaration.project, message, IdeBundle.message("title.warning"), Messages.getQuestionIcon(), Messages.YES)
         when (exitCode) {
             Messages.YES -> return overriddenElementsToDescriptor.keys.toList()
             Messages.NO -> return listOf(declaration)

idea/testData/findUsages/kotlin/conventions/equalsNotAny.0.kt

@@ -1,11 +1,12 @@
 // PSI_ELEMENT: org.jetbrains.kotlin.psi.KtNamedFunction
+// CHECK_SUPER_METHODS_YES_NO_DIALOG: no
 // OPTIONS: usages
 // FIND_BY_REF
 
 class A(val n: Int) {
     fun foo(a: A) = a <caret>== this
 
-    override override fun equals(other: Any?): Boolean = TODO()
+    override infix fun equals(other: Any?): Boolean = TODO()
 }
 
 fun test() {

idea/testData/findUsages/kotlin/conventions/equalsNotAny.results.txt

@@ -1,4 +1,4 @@
-Function call 13 A(0) == A(1)
-Function call 14 A(0) != A(1)
-Function call 15 A(0) equals A(1)
-Function call 6 fun foo(a: A) = a == this
+Function call 14 A(0) == A(1)
+Function call 15 A(0) != A(1)
+Function call 16 A(0) equals A(1)
+Function call 7 fun foo(a: A) = a == this

idea/testData/findUsages/kotlin/findFunctionUsages/highlightingOfSuperUsages.0.kt

@@ -1,4 +1,5 @@
 // PSI_ELEMENT: org.jetbrains.kotlin.psi.KtNamedFunction
+// CHECK_SUPER_METHODS_YES_NO_DIALOG: no
 // OPTIONS: usages, skipImports
 // HIGHLIGHTING
 

idea/testData/findUsages/kotlin/findFunctionUsages/highlightingOfSuperUsages.results.txt

@@ -1,3 +1,3 @@
-Function call 14 i.foo()
-Function call 15 A().foo()
-Function call 16 B().foo()
+Function call 15 i.foo()
+Function call 16 A().foo()
+Function call 17 B().foo()

idea/testData/findUsages/kotlin/findJavaPropertyUsages/javaPropertyUsagesKJK.0.kt

@@ -1,4 +1,5 @@
 // PSI_ELEMENT: org.jetbrains.kotlin.psi.KtProperty
+// CHECK_SUPER_METHODS_YES_NO_DIALOG: no
 // OPTIONS: usages
 open class A {
     open var p: Int = 1

idea/testData/findUsages/kotlin/findJavaPropertyUsages/javaPropertyUsagesKJK.results.txt

@@ -1,4 +1,4 @@
-[javaPropertyUsagesKJK.0.kt] Value read 25 val t = B().p
-[javaPropertyUsagesKJK.0.kt] Value write 26 B().p = 1
+[javaPropertyUsagesKJK.0.kt] Value read 26 val t = B().p
+[javaPropertyUsagesKJK.0.kt] Value write 27 B().p = 1
 [javaPropertyUsagesKJK.1.java] Unclassified usage 24 new B().getP();
-[javaPropertyUsagesKJK.1.java] Unclassified usage 25 new B().setP(1);
+[javaPropertyUsagesKJK.1.java] Unclassified usage 25 new B().setP(1);

idea/testData/findUsages/kotlin/findJavaPropertyUsages/javaPropertyUsagesKK.0.kt

@@ -1,4 +1,5 @@
 // PSI_ELEMENT: org.jetbrains.kotlin.psi.KtProperty
+// CHECK_SUPER_METHODS_YES_NO_DIALOG: no
 // OPTIONS: usages
 open class A {
     open var p: Int = 1

idea/testData/findUsages/kotlin/findJavaPropertyUsages/javaPropertyUsagesKK.results.txt

@@ -1,4 +1,4 @@
-[javaPropertyUsagesKK.0.kt] Value read 19 val t = AA().p
-[javaPropertyUsagesKK.0.kt] Value write 20 AA().p = 1
+[javaPropertyUsagesKK.0.kt] Value read 20 val t = AA().p
+[javaPropertyUsagesKK.0.kt] Value write 21 AA().p = 1
 [javaPropertyUsagesKK.1.java] Unclassified usage 18 new AA().getP();
-[javaPropertyUsagesKK.1.java] Unclassified usage 19 new AA().setP(1);
+[javaPropertyUsagesKK.1.java] Unclassified usage 19 new AA().setP(1);

idea/tests/org/jetbrains/kotlin/findUsages/AbstractFindUsagesTest.kt

@@ -29,6 +29,7 @@ import com.intellij.lang.properties.psi.Property
 import com.intellij.openapi.editor.markup.TextAttributes
 import com.intellij.openapi.extensions.Extensions
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.ui.Messages
 import com.intellij.openapi.util.io.FileUtil
 import com.intellij.openapi.util.io.FileUtilRt
 import com.intellij.psi.PsiComment
@@ -46,14 +47,14 @@ import com.intellij.usages.impl.rules.UsageTypeProvider
 import com.intellij.usages.rules.UsageFilteringRule
 import com.intellij.usages.rules.UsageGroupingRule
 import com.intellij.util.CommonProcessors
-import org.jetbrains.kotlin.idea.findUsages.KotlinFindUsagesHandlerFactory
+import org.jetbrains.kotlin.idea.core.util.clearDialogsResults
+import org.jetbrains.kotlin.idea.core.util.setDialogsResult
+import org.jetbrains.kotlin.idea.refactoring.CHECK_SUPER_METHODS_YES_NO_DIALOG
 import org.jetbrains.kotlin.idea.search.usagesSearch.ExpressionsOfTypeProcessor
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
 import org.jetbrains.kotlin.idea.test.TestFixtureExtension
 import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
-import org.jetbrains.kotlin.psi.KtDeclaration
-import org.jetbrains.kotlin.psi.KtTypeAlias
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
 import org.jetbrains.kotlin.test.KotlinTestUtils
@@ -206,7 +207,10 @@ internal fun <T : PsiElement> findUsagesAndCheckResults(
             ExpressionsOfTypeProcessor.mode = ExpressionsOfTypeProcessor.Mode.PLAIN_WHEN_NEEDED
         }
 
-        findUsages(caretElement, options, highlightingMode, project)
+        val searchSuperDeclaration =
+                InTextDirectivesUtils.findLinesWithPrefixesRemoved(mainFileText, CHECK_SUPER_METHODS_YES_NO_DIALOG + ":").firstOrNull() != "no"
+
+        findUsages(caretElement, options, highlightingMode, project, searchSuperDeclaration)
     }
     finally {
         ExpressionsOfTypeProcessor.testLog = null
@@ -269,46 +273,54 @@ internal fun <T : PsiElement> findUsagesAndCheckResults(
     }
 }
 
-
 internal fun findUsages(
         targetElement: PsiElement,
         options: FindUsagesOptions?,
         highlightingMode: Boolean,
-        project: Project
+        project: Project,
+        searchSuperDeclaration: Boolean = true
 ): Collection<UsageInfo> {
-    val handler: FindUsagesHandler = (if (targetElement is PsiMember) {
-        JavaFindUsagesHandler(targetElement, JavaFindUsagesHandlerFactory(project))
-    }
-    else if (targetElement is KtDeclaration && targetElement !is KtTypeAlias) {
-        KotlinFindUsagesHandlerFactory(project).createFindUsagesHandlerNoQuestions(targetElement)
-    }
-    else {
-        (FindManager.getInstance(project) as FindManagerImpl).findUsagesManager.getFindUsagesHandler(targetElement, false)
-    }) ?: error("Cannot find handler for: $targetElement")
+    try {
+        val handler: FindUsagesHandler = when {
+            targetElement is PsiMember ->
+                JavaFindUsagesHandler(targetElement, JavaFindUsagesHandlerFactory(project))
+            else -> {
+                if (!searchSuperDeclaration) {
+                    setDialogsResult(CHECK_SUPER_METHODS_YES_NO_DIALOG, Messages.NO)
+                }
 
-    @Suppress("NAME_SHADOWING")
-    val options = options ?: handler.getFindUsagesOptions(null)
-
-    options.searchScope = GlobalSearchScope.allScope(project)
-
-    val processor = CommonProcessors.CollectProcessor<UsageInfo>()
-    for (psiElement in handler.primaryElements + handler.secondaryElements) {
-        if (highlightingMode) {
-            //TODO: should findReferencesToHighlight work outside read-action or it makes no sense?
-            for (reference in handler.findReferencesToHighlight(psiElement, options.searchScope)) {
-                processor.process(UsageInfo(reference))
+                val findManagerImpl = FindManager.getInstance(project) as FindManagerImpl
+                findManagerImpl.findUsagesManager.getFindUsagesHandler(targetElement, false) ?: error("Cannot find handler for: $targetElement")
             }
         }
-        else {
-            // run in another thread to test read-action assertions
-            val thread = Thread {
-                handler.processElementUsages(psiElement, processor, options)
-            }
-            thread.start()
-            thread.join()
-        }
-    }
 
-    return processor.results
+        @Suppress("NAME_SHADOWING")
+        val options = options ?: handler.getFindUsagesOptions(null)
+
+        options.searchScope = GlobalSearchScope.allScope(project)
+
+        val processor = CommonProcessors.CollectProcessor<UsageInfo>()
+        for (psiElement in handler.primaryElements + handler.secondaryElements) {
+            if (highlightingMode) {
+                //TODO: should findReferencesToHighlight work outside read-action or it makes no sense?
+                for (reference in handler.findReferencesToHighlight(psiElement, options.searchScope)) {
+                    processor.process(UsageInfo(reference))
+                }
+            }
+            else {
+                // run in another thread to test read-action assertions
+                val thread = Thread {
+                    handler.processElementUsages(psiElement, processor, options)
+                }
+                thread.start()
+                thread.join()
+            }
+        }
+
+        return processor.results
+    }
+    finally {
+        clearDialogsResults()
+    }
 }