Add kotlin-test.jar to classpath quickfix

idea/idea-analysis/src/org/jetbrains/kotlin/idea/KotlinBundle.properties

@@ -263,6 +263,7 @@ suppress.warning.for=Suppress ''{0}'' for {1} {2}
 
 reflection.not.found=Reflection not found in the classpath
 add.reflection.to.classpath=Add kotlin-reflect.jar to the classpath
+add.test.to.classpath=Add kotlin-test.jar to the classpath
 
 # Kotlin Compiler Settings Tab
 

idea/src/org/jetbrains/kotlin/idea/framework/JavaRuntimePresentationProvider.java

@@ -72,4 +72,9 @@ public class JavaRuntimePresentationProvider extends LibraryPresentationProvider
     public static VirtualFile getRuntimeSrcJar(@NotNull Library library) {
         return getRuntimeSrcJar(Arrays.asList(library.getFiles(OrderRootType.SOURCES)));
     }
+
+    @Nullable
+    public static VirtualFile getTestJar(@NotNull Library library) {
+        return LibraryUtils.getJarFile(Arrays.asList(library.getFiles(OrderRootType.CLASSES)), PathUtil.KOTLIN_TEST_JAR);
+    }
 }

idea/src/org/jetbrains/kotlin/idea/inspections/AddKotlinLibQuickFix.kt

@@ -16,13 +16,16 @@
 
 package org.jetbrains.kotlin.idea.inspections
 
+import com.intellij.codeInsight.intention.IntentionAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.extensions.Extensions
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.OrderRootType
+import com.intellij.openapi.roots.libraries.Library
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VfsUtilCore
 import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.configuration.ConfigureKotlinInProjectUtils
 import org.jetbrains.kotlin.idea.configuration.KotlinJavaModuleConfigurator
@@ -38,42 +41,80 @@ import org.jetbrains.kotlin.utils.PathUtil
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import java.io.File
 
-public class AddReflectionQuickFix(element: KtElement) : KotlinQuickFixAction<KtElement>(element) {
+public class AddReflectionQuickFix(element: KtElement) : AddKotlinLibQuickFix(element) {
     override fun getText() = KotlinBundle.message("add.reflection.to.classpath")
-    override fun getFamilyName() = getText()
+    override fun getFamilyName() = text
+
+    override fun libraryPath(): String = PathUtil.KOTLIN_JAVA_REFLECT_JAR
+    override fun getLibFile(): File = PathUtil.getKotlinPathsForIdeaPlugin().getReflectPath()
+    override fun hasLibJarInLibrary(library: Library): Boolean = JavaRuntimePresentationProvider.getReflectJar(library) != null
+
+    companion object : KotlinSingleIntentionActionFactory() {
+        override fun createAction(diagnostic: Diagnostic) = diagnostic.createIntentionForFirstParentOfType(::AddReflectionQuickFix)
+    }
+}
+
+public class AddTestLibQuickFix(element: KtElement) : AddKotlinLibQuickFix(element) {
+    override fun getText() = KotlinBundle.message("add.test.to.classpath")
+    override fun getFamilyName() = text
+
+    override fun libraryPath(): String = PathUtil.KOTLIN_TEST_JAR
+    override fun getLibFile(): File = PathUtil.getKotlinPathsForIdeaPlugin().kotlinTestPath
+    override fun hasLibJarInLibrary(library: Library): Boolean = JavaRuntimePresentationProvider.getTestJar(library) != null
+
+    companion object : KotlinSingleIntentionActionFactory() {
+        val KOTLIN_TEST_UNRESOLVED = setOf(
+                "Asserter", "assertFailsWith", "currentStackTrace", "failsWith", "todo", "assertEquals",
+                "assertFails", "assertNot", "assertNotEquals", "assertNotNull", "assertNull", "assertTrue", "expect", "fail", "fails")
+
+        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
+
+                return diagnostic.createIntentionForFirstParentOfType(::AddTestLibQuickFix)
+            }
+
+            return null
+        }
+    }
+}
+
+public abstract class AddKotlinLibQuickFix(element: KtElement) : KotlinQuickFixAction<KtElement>(element) {
+    protected abstract fun libraryPath(): String
+    protected abstract fun getLibFile(): File
+    protected abstract fun hasLibJarInLibrary(library: Library): Boolean
 
     override fun invoke(project: Project, editor: Editor?, file: KtFile) {
-        val pluginReflectJar = PathUtil.getKotlinPathsForIdeaPlugin().getReflectPath()
-        if (!pluginReflectJar.exists()) return
+        val libFile = getLibFile()
+        if (!libFile.exists()) return
 
         val configurator = Extensions.getExtensions(KotlinProjectConfigurator.EP_NAME)
                                    .firstIsInstanceOrNull<KotlinJavaModuleConfigurator>() ?: return
 
         for (library in KotlinRuntimeLibraryUtil.findKotlinLibraries(project)) {
             val runtimeJar = JavaRuntimePresentationProvider.getRuntimeJar(library) ?: continue
-            if (JavaRuntimePresentationProvider.getReflectJar(library) != null) continue
+            if (hasLibJarInLibrary(library)) continue
 
-            val model = library.getModifiableModel()
+            val model = library.modifiableModel
 
-            val libFilesDir = VfsUtilCore.virtualToIoFile(runtimeJar).getParent()
-            val reflectIoFile = File(libFilesDir, PathUtil.KOTLIN_JAVA_REFLECT_JAR)
-            if (reflectIoFile.exists()) {
-                model.addRoot(VfsUtil.getUrlForLibraryRoot(reflectIoFile), OrderRootType.CLASSES)
+            val libFilesDir = VfsUtilCore.virtualToIoFile(runtimeJar).parent
+
+            val libIoFile = File(libFilesDir, libraryPath())
+            if (libIoFile.exists()) {
+                model.addRoot(VfsUtil.getUrlForLibraryRoot(libIoFile), OrderRootType.CLASSES)
             }
             else {
-                val copied = configurator.copyFileToDir(project, pluginReflectJar, libFilesDir)!!
+                val copied = configurator.copyFileToDir(project, libFile, libFilesDir)!!
                 model.addRoot(VfsUtil.getUrlForLibraryRoot(copied), OrderRootType.CLASSES)
             }
 
             model.commit()
 
-            ConfigureKotlinInProjectUtils.showInfoNotification(project,
-                    "${PathUtil.KOTLIN_JAVA_REFLECT_JAR} was added to the library ${library.getName()}"
-            )
+            ConfigureKotlinInProjectUtils.showInfoNotification(
+                    project, "${libraryPath()} was added to the library ${library.name}")
         }
     }
-
-    companion object : KotlinSingleIntentionActionFactory() {
-        override fun createAction(diagnostic: Diagnostic) = diagnostic.createIntentionForFirstParentOfType(::AddReflectionQuickFix)
-    }
 }

idea/src/org/jetbrains/kotlin/idea/quickfix/QuickFixRegistrar.kt

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.diagnostics.Errors.*
 import org.jetbrains.kotlin.idea.core.overrideImplement.ImplementMembersHandler
 import org.jetbrains.kotlin.idea.inspections.AddModifierFixFactory
 import org.jetbrains.kotlin.idea.inspections.AddReflectionQuickFix
+import org.jetbrains.kotlin.idea.inspections.AddTestLibQuickFix
 import org.jetbrains.kotlin.idea.inspections.InfixCallFix
 import org.jetbrains.kotlin.idea.quickfix.createFromUsage.createCallable.*
 import org.jetbrains.kotlin.idea.quickfix.createFromUsage.createClass.CreateClassFromCallWithConstructorCalleeActionFactory
@@ -125,6 +126,8 @@ public class QuickFixRegistrar : QuickFixContributor {
         REPEATED_MODIFIER.registerFactory(removeModifierFactory)
 
         UNRESOLVED_REFERENCE.registerFactory(AutoImportFix)
+        UNRESOLVED_REFERENCE.registerFactory(AddTestLibQuickFix)
+
         UNRESOLVED_REFERENCE_WRONG_RECEIVER.registerFactory(AutoImportFix)
 
         FUNCTION_EXPECTED.registerFactory(MissingInvokeAutoImportFix)