From f25d80c38b43095e725c54902e703c91e1ad2a0f Mon Sep 17 00:00:00 2001
From: Nikolay Krasko <goodwinnk@gmail.com>
Date: Tue, 22 Dec 2015 15:57:11 +0300
Subject: [PATCH] Make add test lib quick-fix applicable only when there is
 unresolved import

---
 .../idea/inspections/AddKotlinLibQuickFix.kt  | 34 ++++++++++++++++---
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/inspections/AddKotlinLibQuickFix.kt b/idea/src/org/jetbrains/kotlin/idea/inspections/AddKotlinLibQuickFix.kt
index 6ddbf59021c..490978eccde 100644
--- a/idea/src/org/jetbrains/kotlin/idea/inspections/AddKotlinLibQuickFix.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/inspections/AddKotlinLibQuickFix.kt
@@ -30,6 +30,7 @@ import com.intellij.openapi.roots.libraries.Library
 import com.intellij.openapi.roots.ui.configuration.projectRoot.LibrariesContainerFactory
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VfsUtilCore
+import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.idea.KotlinBundle
@@ -43,8 +44,10 @@ import org.jetbrains.kotlin.idea.quickfix.KotlinQuickFixAction
 import org.jetbrains.kotlin.idea.quickfix.KotlinSingleIntentionActionFactory
 import org.jetbrains.kotlin.idea.quickfix.quickfixUtil.createIntentionForFirstParentOfType
 import org.jetbrains.kotlin.idea.versions.KotlinRuntimeLibraryUtil
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtImportDirective
 import org.jetbrains.kotlin.utils.PathUtil
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import java.io.File
@@ -82,11 +85,34 @@ public class AddTestLibQuickFix(element: KtElement) : AddKotlinLibQuickFix(eleme
         override fun createAction(diagnostic: Diagnostic): IntentionAction? {
             val unresolvedReference = Errors.UNRESOLVED_REFERENCE.cast(diagnostic)
 
-            if (unresolvedReference.a.text in KOTLIN_TEST_UNRESOLVED) {
-                val ktFile = (diagnostic.psiElement.containingFile as? KtFile) ?: return null
-                if (ktFile.importDirectives.none { it.text.contains("kotlin.test.") }) return null
+            if (PsiTreeUtil.getParentOfType(diagnostic.psiElement, KtImportDirective::class.java) != null) return null
+
+            val unresolvedText = unresolvedReference.a.text
+            if (unresolvedText in KOTLIN_TEST_UNRESOLVED) {
+                val ktFile = (diagnostic.psiElement.containingFile as? KtFile) ?: return null
+
+                val exactImportFqName = FqName("kotlin.test.$unresolvedText")
+                val kotlinTestAllUnder = FqName("kotlin.test")
+
+                var hasExactImport = false
+                var hasKotlinTestAllUnder = false
+
+                for (importDirective in ktFile.importDirectives.filter { it.text.contains("kotlin.test.") }) {
+                    if (importDirective.importedFqName == exactImportFqName) {
+                        hasExactImport = true
+                        break
+                    }
+
+                    if (importDirective.importedFqName == kotlinTestAllUnder && importDirective.isAllUnder) {
+                        hasKotlinTestAllUnder = true
+                        break
+                    }
+                }
+
+                if (hasExactImport || hasKotlinTestAllUnder) {
+                    return diagnostic.createIntentionForFirstParentOfType(::AddTestLibQuickFix)
+                }
 
-                return diagnostic.createIntentionForFirstParentOfType(::AddTestLibQuickFix)
             }
 
             return null