From 0d1a89d8121cbd2ff664c45cffdf1353a8df0af8 Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Tue, 11 Jan 2022 19:08:10 +0100
Subject: [PATCH] [kotlin] [tests] further isConfiguredBy optimization
 (IDEA-285422)

and kotlin+testNG simplification

GitOrigin-RevId: d59fc012c9315b8c89c0606754886643edae3ec3
---
 ...AbstractJavaTestConfigurationProducer.java | 18 ++++---
 .../KotlinTestNgConfigurationProducer.java    | 47 ++++---------------
 2 files changed, 20 insertions(+), 45 deletions(-)

diff --git a/java/execution/impl/src/com/intellij/execution/testframework/AbstractJavaTestConfigurationProducer.java b/java/execution/impl/src/com/intellij/execution/testframework/AbstractJavaTestConfigurationProducer.java
index ef02cc7e9403..5effcd496c21 100644
--- a/java/execution/impl/src/com/intellij/execution/testframework/AbstractJavaTestConfigurationProducer.java
+++ b/java/execution/impl/src/com/intellij/execution/testframework/AbstractJavaTestConfigurationProducer.java
@@ -1,4 +1,4 @@
-// Copyright 2000-2021 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+// Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.execution.testframework;
 
 import com.intellij.codeInsight.TestFrameworks;
@@ -86,6 +86,7 @@ public abstract class AbstractJavaTestConfigurationProducer<T extends JavaTestCo
     if (isMultipleElementsSelected(context)) {
       return false;
     }
+    if (!isApplicableTestType(configuration.getTestType(), context)) return false;
     final RunConfiguration predefinedConfiguration = context.getOriginalConfiguration(getConfigurationType());
     final Location contextLocation = context.getLocation();
     if (contextLocation == null) {
@@ -111,12 +112,9 @@ public abstract class AbstractJavaTestConfigurationProducer<T extends JavaTestCo
     if (!Comparing.strEqual(vmParameters, configuration.getVMParameters())) return false;
     if (differentParamSet(configuration, contextLocation)) return false;
 
-    if (!isApplicableTestType(configuration.getTestType(), context)) return false;
+    PsiElement psiElement = getElement(element);
 
-    PsiClass psiClass = PsiTreeUtil.getParentOfType(element, PsiClass.class);
-    if (psiClass != null && getCurrentFramework(psiClass) == null) return false;
-
-    if (configuration.isConfiguredByElement(element)) {
+    if (psiElement != null && configuration.isConfiguredByElement(psiElement)) {
       final Module configurationModule = configuration.getConfigurationModule().getModule();
       final Module locationModule = location.getModule();
       if (Comparing.equal(locationModule, configurationModule)) return true;
@@ -126,6 +124,14 @@ public abstract class AbstractJavaTestConfigurationProducer<T extends JavaTestCo
     return false;
   }
 
+  protected PsiElement getElement(PsiElement element) {
+    PsiClass psiClass = PsiTreeUtil.getParentOfType(element, PsiClass.class);
+    if (psiClass != null && getCurrentFramework(psiClass) == null) {
+      return null;
+    }
+    return element;
+  }
+
   protected boolean differentParamSet(T configuration, Location contextLocation) {
     String paramSetName = contextLocation instanceof PsiMemberParameterizedLocation
                           ? configuration.prepareParameterizedParameter(((PsiMemberParameterizedLocation)contextLocation).getParamSetName()) : null;
diff --git a/plugins/kotlin/testng/src/org/jetbrains/kotlin/idea/testng/KotlinTestNgConfigurationProducer.java b/plugins/kotlin/testng/src/org/jetbrains/kotlin/idea/testng/KotlinTestNgConfigurationProducer.java
index c7606cbc21fa..25d5e16f68a4 100644
--- a/plugins/kotlin/testng/src/org/jetbrains/kotlin/idea/testng/KotlinTestNgConfigurationProducer.java
+++ b/plugins/kotlin/testng/src/org/jetbrains/kotlin/idea/testng/KotlinTestNgConfigurationProducer.java
@@ -2,15 +2,16 @@
 
 package org.jetbrains.kotlin.idea.testng;
 
-import com.intellij.execution.*;
+import com.intellij.execution.JavaRunConfigurationExtensionManager;
+import com.intellij.execution.Location;
+import com.intellij.execution.PsiLocation;
+import com.intellij.execution.RunnerAndConfigurationSettings;
 import com.intellij.execution.actions.ConfigurationContext;
 import com.intellij.execution.actions.ConfigurationFromContext;
 import com.intellij.execution.configurations.ConfigurationFactory;
-import com.intellij.execution.configurations.RunConfiguration;
 import com.intellij.execution.testframework.AbstractInClassConfigurationProducer;
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.project.Project;
-import com.intellij.openapi.util.Comparing;
 import com.intellij.openapi.util.Ref;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiElement;
@@ -35,42 +36,10 @@ public class KotlinTestNgConfigurationProducer extends TestNGConfigurationProduc
     }
 
     @Override
-    public boolean isConfigurationFromContext(@NotNull TestNGConfiguration configuration, @NotNull ConfigurationContext context) {
-        if (isMultipleElementsSelected(context)) {
-            return false;
-        }
-        final RunConfiguration predefinedConfiguration = context.getOriginalConfiguration(getConfigurationType());
-        final Location contextLocation = context.getLocation();
-        if (contextLocation == null) {
-            return false;
-        }
-        Location location = JavaExecutionUtil.stepIntoSingleClass(contextLocation);
-        if (location == null) {
-            return false;
-        }
-        final PsiElement element = location.getPsiElement();
-
-        RunnerAndConfigurationSettings template =
-                RunManager.getInstance(location.getProject()).getConfigurationTemplate(getConfigurationFactory());
-        final Module predefinedModule = ((TestNGConfiguration) template.getConfiguration()).getConfigurationModule().getModule();
-        final String vmParameters =
-                predefinedConfiguration instanceof CommonJavaRunConfigurationParameters
-                ? ((CommonJavaRunConfigurationParameters) predefinedConfiguration).getVMParameters()
-                : null;
-        if (vmParameters != null && !Comparing.strEqual(vmParameters, configuration.getVMParameters())) return false;
-        if (differentParamSet(configuration, contextLocation)) return false;
-
-        KotlinTestFrameworkProvider.JavaEntity testEntity = TestNgKotlinTestFrameworkProvider.INSTANCE.getJavaEntity(element);
-        if (testEntity == null) return false;
-        PsiNamedElement lightElement = testEntity.getMethod() != null ? testEntity.getMethod() 
-                                                                      : testEntity.getTestClass();
-        if (lightElement != null && configuration.isConfiguredByElement(lightElement)) {
-            final Module configurationModule = configuration.getConfigurationModule().getModule();
-            if (Comparing.equal(location.getModule(), configurationModule)) return true;
-            if (Comparing.equal(predefinedModule, configurationModule)) return true;
-        }
-
-        return false;
+    protected PsiElement getElement(PsiElement element) {
+         KotlinTestFrameworkProvider.JavaEntity testEntity = TestNgKotlinTestFrameworkProvider.INSTANCE.getJavaEntity(element);
+        if (testEntity == null) return null;
+        return testEntity.getMethod() != null ? testEntity.getMethod() : testEntity.getTestClass();
     }
 
     @Override