Line Marking: Implement run markers for test classes/functions

idea/idea-test-framework/src/org/jetbrains/kotlin/idea/test/ConfigLibraryUtil.java

@@ -29,9 +29,13 @@ import kotlin.jvm.functions.Function0;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.idea.js.KotlinJavaScriptLibraryManager;
 import org.jetbrains.kotlin.idea.util.application.ApplicationUtilsKt;
+import org.jetbrains.kotlin.test.InTextDirectivesUtils;
 import org.jetbrains.kotlin.utils.PathUtil;
 
 import java.io.File;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
 
 /**
  * Helper for configuring kotlin runtime in tested project.
@@ -168,4 +172,39 @@ public class ConfigLibraryUtil {
                 }
         );
     }
+
+    public static void configureLibrariesByDirective(@NotNull Module module, String rootPath, String fileText) {
+        for (String libraryInfo : InTextDirectivesUtils.findListWithPrefixes(fileText, "// CONFIGURE_LIBRARY: ")) {
+            int i = libraryInfo.indexOf('@');
+            String libraryName = libraryInfo.substring(0, i);
+            String[] jarPaths = libraryInfo.substring(i + 1).split(";");
+
+            NewLibraryEditor editor = new NewLibraryEditor();
+            editor.setName(libraryName);
+            for (String jarPath : jarPaths) {
+                editor.addRoot(VfsUtil.getUrlForLibraryRoot(new File(rootPath, jarPath)), OrderRootType.CLASSES);
+            }
+
+            addLibrary(editor, module);
+        }
+    }
+
+    public static void unconfigureLibrariesByDirective(@NotNull Module module, String fileText) {
+        List<String> libraryNames = new ArrayList<String>();
+        for (String libInfo : InTextDirectivesUtils.findListWithPrefixes(fileText, "// CONFIGURE_LIBRARY: ")) {
+            libraryNames.add(libInfo.substring(0, libInfo.indexOf('@')));
+        }
+        for (String libraryName : InTextDirectivesUtils.findListWithPrefixes(fileText, "// UNCONFIGURE_LIBRARY: ")) {
+            libraryNames.add(libraryName);
+        }
+
+        for (Iterator<String> iterator = libraryNames.iterator(); iterator.hasNext(); ) {
+            String libraryName = iterator.next();
+            if (removeLibrary(module, libraryName)) {
+                iterator.remove();
+            }
+        }
+
+        if (!libraryNames.isEmpty()) throw new AssertionError("Couldn't find the following libraries: " + libraryNames);
+    }
 }

idea/src/META-INF/plugin.xml

@@ -443,6 +443,7 @@
     <codeInsight.lineMarkerProvider language="kotlin" implementationClass="org.jetbrains.kotlin.idea.highlighter.markers.KotlinLineMarkerProvider"/>
     <codeInsight.lineMarkerProvider language="kotlin" implementationClass="org.jetbrains.kotlin.idea.highlighter.KotlinRecursiveCallLineMarkerProvider"/>
     <runLineMarkerContributor language="kotlin" implementationClass="org.jetbrains.kotlin.idea.highlighter.KotlinRunLineMarkerContributor"/>
+    <runLineMarkerContributor language="kotlin" implementationClass="org.jetbrains.kotlin.idea.highlighter.KotlinTestRunLineMarkerContributor"/>
 
     <iconProvider implementation="org.jetbrains.kotlin.idea.JetIconProvider"/>
     <itemPresentationProvider implementationClass="org.jetbrains.kotlin.idea.presentation.JetFunctionPresenter"

idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinTestRunLineMarkerContributor.kt

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2010-2015 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.highlighter
+
+import com.intellij.codeInsight.TestFrameworks
+import com.intellij.execution.TestStateStorage
+import com.intellij.execution.lineMarker.ExecutorAction
+import com.intellij.execution.lineMarker.RunLineMarkerContributor
+import com.intellij.execution.testframework.TestIconMapper
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.asJava.KotlinLightClass
+import org.jetbrains.kotlin.asJava.toLightClass
+import org.jetbrains.kotlin.asJava.toLightMethods
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.psi.JetClassOrObject
+import org.jetbrains.kotlin.psi.JetNamedDeclaration
+import org.jetbrains.kotlin.psi.JetNamedFunction
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import javax.swing.Icon
+
+class KotlinTestRunLineMarkerContributor : RunLineMarkerContributor() {
+    private fun getTestStateIcon(url: String, project: Project): Icon? {
+        val state = TestStateStorage.getInstance(project).getState(url) ?: return null
+        val magnitude = TestIconMapper.getMagnitude(state.magnitude)
+        return TestIconMapper.getIcon(magnitude)
+    }
+
+    override fun getInfo(element: PsiElement): RunLineMarkerContributor.Info? {
+        val declaration = element.getStrictParentOfType<JetNamedDeclaration>() ?: return null
+        if (declaration.nameIdentifier != element) return null
+
+        // To prevent IDEA failing on red code
+        if (declaration.analyze(BodyResolveMode.PARTIAL)
+                .get(BindingContext.DECLARATION_TO_DESCRIPTOR, declaration) == null) return null
+
+        val project = element.project
+
+        val (url, framework) = when (declaration) {
+            is JetClassOrObject -> {
+                val lightClass = declaration.toLightClass() ?: return null
+                val framework = TestFrameworks.detectFramework(lightClass) ?: return null
+                if (!framework.isTestClass(lightClass)) return null
+
+                "java:suite://${lightClass.qualifiedName!!}" to framework
+            }
+
+            is JetNamedFunction -> {
+                val lightMethod = declaration.toLightMethods().firstOrNull() ?: return null
+                val lightClass = lightMethod.containingClass as? KotlinLightClass ?: return null
+                val framework = TestFrameworks.detectFramework(lightClass) ?: return null
+                if (!framework.isTestMethod(lightMethod)) return null
+
+                "java:test://${lightClass.qualifiedName}.${lightMethod.name}" to framework
+            }
+
+            else -> return null
+        }
+
+        val icon = getTestStateIcon(url, project) ?: framework.icon
+        return RunLineMarkerContributor.Info(icon, { "Run Test" }, ExecutorAction.getActions(1))
+    }
+}

idea/src/org/jetbrains/kotlin/idea/run/JetTestNgConfigurationProducer.java

@@ -20,23 +20,30 @@ import com.intellij.execution.JavaRunConfigurationExtensionManager;
 import com.intellij.execution.Location;
 import com.intellij.execution.PsiLocation;
 import com.intellij.execution.actions.ConfigurationContext;
+import com.intellij.execution.actions.ConfigurationFromContext;
+import com.intellij.execution.junit.InheritorChooser;
+import com.intellij.execution.junit2.info.MethodLocation;
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.Condition;
 import com.intellij.openapi.util.Ref;
-import com.intellij.psi.PsiClass;
+import com.intellij.psi.*;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiMethod;
 import com.intellij.psi.util.PsiClassUtil;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.theoryinpractice.testng.configuration.TestNGConfiguration;
 import com.theoryinpractice.testng.configuration.TestNGConfigurationProducer;
 import com.theoryinpractice.testng.util.TestNGUtil;
+import kotlin.CollectionsKt;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.asJava.LightClassUtil;
+import org.jetbrains.kotlin.asJava.LightClassUtilsKt;
 import org.jetbrains.kotlin.idea.project.ProjectStructureUtil;
 import org.jetbrains.kotlin.idea.util.ProjectRootsUtil;
 import org.jetbrains.kotlin.psi.*;
 
+import java.util.List;
+
 public class JetTestNgConfigurationProducer extends TestNGConfigurationProducer {
     @Override
     protected boolean setupConfigurationFromContext(
@@ -67,8 +74,11 @@ public class JetTestNgConfigurationProducer extends TestNGConfigurationProducer
             return false;
         }
 
-        JetNamedFunction function = PsiTreeUtil.getParentOfType(leaf, JetNamedFunction.class, false);
+        JetNamedDeclaration declarationToRun = getDeclarationToRun(leaf);
-        if (function != null) {
+
+        if (declarationToRun instanceof JetNamedFunction) {
+            JetNamedFunction function = (JetNamedFunction) declarationToRun;
+
             @SuppressWarnings("unchecked")
             JetElement owner = PsiTreeUtil.getParentOfType(function, JetFunction.class, JetClass.class);
 
@@ -87,22 +97,86 @@ public class JetTestNgConfigurationProducer extends TestNGConfigurationProducer
             }
         }
 
+        if (declarationToRun instanceof JetClass) {
+            PsiClass delegate = LightClassUtil.INSTANCE$.getPsiClass((JetClassOrObject) declarationToRun);
+            if (!isTestNGClass(delegate)) {
+                return false;
+            }
+
+            return configure(configuration, location, context, project, delegate, null);
+        }
+
+        return false;
+    }
+
+    @Override
+    public void onFirstRun(ConfigurationFromContext configuration, ConfigurationContext context, Runnable startRunnable) {
+        JetNamedDeclaration declarationToRun = getDeclarationToRun(configuration.getSourceElement());
+        final PsiNamedElement lightElement = CollectionsKt.firstOrNull(LightClassUtilsKt.toLightElements(declarationToRun));
+
+        // Copied from TestNGInClassConfigurationProducer.onFirstRun()
+        if (lightElement instanceof PsiMethod || lightElement instanceof PsiClass) {
+            PsiMethod psiMethod;
+            PsiClass containingClass;
+
+            if (lightElement instanceof PsiMethod) {
+                psiMethod = (PsiMethod)lightElement;
+                containingClass = psiMethod.getContainingClass();
+            } else {
+                psiMethod = null;
+                containingClass = (PsiClass)lightElement;
+            }
+
+            InheritorChooser inheritorChooser = new InheritorChooser() {
+                @Override
+                protected void runForClasses(List<PsiClass> classes, PsiMethod method, ConfigurationContext context, Runnable performRunnable) {
+                    ((TestNGConfiguration)context.getConfiguration().getConfiguration()).bePatternConfiguration(classes, method);
+                    super.runForClasses(classes, method, context, performRunnable);
+                }
+
+                @Override
+                protected void runForClass(PsiClass aClass,
+                        PsiMethod psiMethod,
+                        ConfigurationContext context,
+                        Runnable performRunnable) {
+                    if (lightElement instanceof PsiMethod) {
+                        Project project = psiMethod.getProject();
+                        MethodLocation methodLocation = new MethodLocation(project, psiMethod, PsiLocation.fromPsiElement(aClass));
+                        ((TestNGConfiguration)context.getConfiguration().getConfiguration()).setMethodConfiguration(methodLocation);
+                    } else {
+                        ((TestNGConfiguration)context.getConfiguration().getConfiguration()).setClassConfiguration(aClass);
+                    }
+                    super.runForClass(aClass, psiMethod, context, performRunnable);
+                }
+            };
+            if (inheritorChooser.runMethodInAbstractClass(context,
+                                                          startRunnable,
+                                                          psiMethod,
+                                                          containingClass,
+                                                          new Condition<PsiClass>() {
+                                                              @Override
+                                                              public boolean value(PsiClass aClass) {
+                                                                  return aClass.hasModifierProperty(PsiModifier.ABSTRACT) &&
+                                                                         TestNGUtil.hasTest(aClass);
+                                                              }
+                                                          })) return;
+        }
+
+        super.onFirstRun(configuration, context, startRunnable);
+    }
+
+    @Nullable
+    private static JetNamedDeclaration getDeclarationToRun(@NotNull PsiElement leaf) {
+        if (!(leaf.getContainingFile() instanceof JetFile)) return null;
+        JetFile jetFile = (JetFile) leaf.getContainingFile();
+
+        JetNamedFunction function = PsiTreeUtil.getParentOfType(leaf, JetNamedFunction.class, false);
+        if (function != null) return function;
+
         JetClass jetClass = PsiTreeUtil.getParentOfType(leaf, JetClass.class, false);
+        if (jetClass != null) return jetClass;
 
-        if (jetClass == null) {
+        return getClassDeclarationInFile(jetFile);
-            jetClass = getClassDeclarationInFile(jetFile);
-        }
-
-        if (jetClass == null) {
-            return false;
-        }
-
-        PsiClass delegate = LightClassUtil.INSTANCE$.getPsiClass(jetClass);
-        if (!isTestNGClass(delegate)) {
-            return false;
-        }
-
-        return configure(configuration, location, context, project, delegate, null);
     }
 
     private boolean configure(

idea/src/org/jetbrains/kotlin/idea/run/KotlinJUnitRunConfigurationProducer.kt

@@ -18,12 +18,10 @@ package org.jetbrains.kotlin.idea.run
 
 import com.intellij.execution.*
 import com.intellij.execution.actions.ConfigurationContext
+import com.intellij.execution.actions.ConfigurationFromContext
 import com.intellij.execution.actions.RunConfigurationProducer
 import com.intellij.execution.configurations.ModuleBasedConfiguration
-import com.intellij.execution.junit.JUnitConfiguration
+import com.intellij.execution.junit.*
-import com.intellij.execution.junit.JUnitConfigurationType
-import com.intellij.execution.junit.JUnitUtil
-import com.intellij.execution.junit.PatternConfigurationProducer
 import com.intellij.openapi.project.DumbService
 import com.intellij.openapi.util.Ref
 import com.intellij.psi.PsiClass
@@ -114,6 +112,26 @@ public class KotlinJUnitRunConfigurationProducer : RunConfigurationProducer<JUni
         return false
     }
 
+    override fun onFirstRun(fromContext: ConfigurationFromContext, context: ConfigurationContext, performRunnable: Runnable) {
+        val leaf = fromContext.sourceElement
+        getTestClass(leaf)?.let { testClass ->
+            val fromContextSubstitute = object : ConfigurationFromContext() {
+                override fun getConfigurationSettings() = fromContext.configurationSettings
+
+                override fun setConfigurationSettings(configurationSettings: RunnerAndConfigurationSettings) {
+                    fromContext.configurationSettings = configurationSettings
+                }
+
+                override fun getSourceElement() = testClass
+            }
+            // TODO: use TestClassConfigurationProducer when constructor becomes public
+            return object : AbstractTestClassConfigurationProducer(JUnitConfigurationType.getInstance()){}
+                    .onFirstRun(fromContextSubstitute, context, performRunnable)
+        }
+
+        super.onFirstRun(fromContext, context, performRunnable)
+    }
+
     private fun getTestMethodLocation(leaf: PsiElement): Location<PsiMethod>? {
         val function = leaf.getParentOfType<JetNamedFunction>(false) ?: return null
         val owner = PsiTreeUtil.getParentOfType(function, javaClass<JetFunction>(), javaClass<JetClass>())
@@ -142,7 +160,7 @@ public class KotlinJUnitRunConfigurationProducer : RunConfigurationProducer<JUni
     }
 
     private fun JetClass?.isJUnitTestClass() =
-            LightClassUtil.getPsiClass(this)?.let { JUnitUtil.isTestClass(it) } ?: false
+            LightClassUtil.getPsiClass(this)?.let { JUnitUtil.isTestClass(it, false, true) } ?: false
 
     private fun getTestClassInFile(jetFile: JetFile) =
             jetFile.getDeclarations().filterIsInstance<JetClass>().singleOrNull { it.isJUnitTestClass() }

idea/testData/codeInsight/lineMarker/runMarkers/jUnitTestClassWithSubclasses.kt

@@ -0,0 +1,26 @@
+// CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+package testing
+
+import junit.framework.TestCase
+import org.junit.Test
+
+abstract class <lineMarker></lineMarker>KBase : TestCase() {
+    @Test
+    fun testFoo() {
+
+    }
+}
+
+class <lineMarker></lineMarker>KTest : KBase() {
+    @Test
+    fun <lineMarker></lineMarker>testBar() {
+
+    }
+}
+
+class <lineMarker></lineMarker>KTest2 : KBase() {
+    @Test
+    fun <lineMarker></lineMarker>testBaz() {
+
+    }
+}

idea/testData/codeInsight/lineMarker/runMarkers/testNGTestClassWithSubclasses.kt

@@ -0,0 +1,25 @@
+// CONFIGURE_LIBRARY: TestNG@plugins/testng/lib/testng.jar
+package testing
+
+import org.testng.annotations.Test
+
+abstract class <lineMarker></lineMarker>KBase {
+    @Test
+    fun <lineMarker></lineMarker>testFoo() {
+
+    }
+}
+
+class <lineMarker></lineMarker>KTest : KBase() {
+    @Test
+    fun <lineMarker></lineMarker>testBar() {
+
+    }
+}
+
+class <lineMarker></lineMarker>KTest2 : KBase() {
+    @Test
+    fun <lineMarker></lineMarker>testBaz() {
+
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/codeInsight/AbstractLineMarkersTest.java

@@ -29,8 +29,10 @@ import com.intellij.psi.PsiElement;
 import com.intellij.rt.execution.junit.FileComparisonFailure;
 import com.intellij.testFramework.ExpectedHighlightingData;
 import com.intellij.testFramework.LightProjectDescriptor;
+import com.intellij.testFramework.PlatformTestUtil;
 import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.kotlin.idea.test.ConfigLibraryUtil;
 import org.jetbrains.kotlin.idea.test.JetLightCodeInsightFixtureTestCase;
 import org.jetbrains.kotlin.idea.test.JetWithJdkAndRuntimeLightProjectDescriptor;
 import org.jetbrains.kotlin.idea.test.PluginTestCaseBase;
@@ -65,6 +67,9 @@ public abstract class AbstractLineMarkersTest extends JetLightCodeInsightFixture
 
     public void doTest(String path) {
         try {
+            String fileText = FileUtil.loadFile(new File(path));
+            ConfigLibraryUtil.configureLibrariesByDirective(myFixture.getModule(), PlatformTestUtil.getCommunityPath(), fileText);
+
             myFixture.configureByFile(path);
             Project project = myFixture.getProject();
             Document document = myFixture.getEditor().getDocument();

idea/tests/org/jetbrains/kotlin/idea/codeInsight/LineMarkersTestGenerated.java

@@ -286,4 +286,25 @@ public class LineMarkersTestGenerated extends AbstractLineMarkersTest {
             doTest(fileName);
         }
     }
+
+    @TestMetadata("idea/testData/codeInsight/lineMarker/runMarkers")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class RunMarkers extends AbstractLineMarkersTest {
+        public void testAllFilesPresentInRunMarkers() throws Exception {
+            JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/codeInsight/lineMarker/runMarkers"), Pattern.compile("^(.+)\\.kt$"), true);
+        }
+
+        @TestMetadata("jUnitTestClassWithSubclasses.kt")
+        public void testJUnitTestClassWithSubclasses() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/codeInsight/lineMarker/runMarkers/jUnitTestClassWithSubclasses.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("testNGTestClassWithSubclasses.kt")
+        public void testTestNGTestClassWithSubclasses() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/codeInsight/lineMarker/runMarkers/testNGTestClassWithSubclasses.kt");
+            doTest(fileName);
+        }
+    }
 }

idea/tests/org/jetbrains/kotlin/idea/quickfix/AbstractQuickFixTest.java

@@ -24,7 +24,6 @@ import com.intellij.codeInspection.SuppressableProblemGroup;
 import com.intellij.ide.startup.impl.StartupManagerImpl;
 import com.intellij.lang.annotation.ProblemGroup;
 import com.intellij.openapi.command.CommandProcessor;
-import com.intellij.openapi.module.Module;
 import com.intellij.openapi.projectRoots.JavaSdk;
 import com.intellij.openapi.projectRoots.Sdk;
 import com.intellij.openapi.startup.StartupManager;
@@ -114,20 +113,12 @@ public abstract class AbstractQuickFixTest extends KotlinLightQuickFixTestCase {
                     e.printStackTrace();
                     fail(testName);
                 } finally {
-                    unconfigureLibrariesAsSpecified(fileText);
+                    ConfigLibraryUtil.unconfigureLibrariesByDirective(getModule(), fileText);
                 }
             }
         }, "", "");
     }
 
-    private static void unconfigureLibrariesAsSpecified(String fileText) {
-        Module module = getModule();
-        for (String libraryName : InTextDirectivesUtils.findListWithPrefixes(fileText, "// UNCONFIGURE_LIBRARY: ")) {
-            if (ConfigLibraryUtil.removeLibrary(module, libraryName)) continue;
-            fail("Library '" + libraryName + "' wasn't found");
-        }
-    }
-
     private static void applyAction(String contents, QuickFixTestCase quickFixTestCase, String testName, String testFullPath) throws Exception {
         Pair<String, Boolean> pair = quickFixTestCase.parseActionHintImpl(quickFixTestCase.getFile(), contents);
         String text = pair.getFirst();