Specify test framework to use explicitly in generate test tests

idea/src/org/jetbrains/kotlin/idea/actions/generate/KotlinGenerateTestSupportActionBase.kt

@@ -127,11 +127,23 @@ abstract class KotlinGenerateTestSupportActionBase(
 
     override fun invoke(project: Project, editor: Editor, file: PsiFile) {
         val klass = findTargetClass(editor, file) ?: return
-        val frameworks = findSuitableFrameworks(klass)
+
-                .filter { methodKind.getFileTemplateDescriptor(it) != null && isApplicableTo(it, klass) }
+        if (testFrameworkToUse != null) {
-        chooseAndPerform(editor, frameworks) { doGenerate(editor, file, klass, it) }
+            val frameworkToUse = findSuitableFrameworks(klass).first { it.name == testFrameworkToUse }
+            if (isApplicableTo(frameworkToUse, klass)) {
+                doGenerate(editor, file, klass, frameworkToUse)
+            }
+        }
+        else {
+            val frameworks = findSuitableFrameworks(klass)
+                    .filter { methodKind.getFileTemplateDescriptor(it) != null && isApplicableTo(it, klass) }
+
+            chooseAndPerform(editor, frameworks) { doGenerate(editor, file, klass, it) }
+        }
     }
 
+    var testFrameworkToUse: String? = null
+
     private val DUMMY_NAME = "__KOTLIN_RULEZZZ__"
 
     private fun doGenerate(editor: Editor, file: PsiFile, klass: KtClassOrObject, framework: TestFramework) {

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/dataMethod.kt

@@ -1,5 +1,6 @@
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$Data
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 class A {<caret>
 
 }

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/dataMethod.kt.after

@@ -2,6 +2,7 @@ import org.junit.runners.Parameterized
 
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$Data
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 class A {
     @Parameterized.Parameters
     fun data(): Collection<Array<Any>> {

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUp.kt

@@ -1,5 +1,6 @@
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$SetUp
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 class A {<caret>
 
 }

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUp.kt.after

@@ -2,6 +2,7 @@ import org.junit.Before
 
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$SetUp
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 class A {
     @Before
     fun setUp() {

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpExists.kt

@@ -1,5 +1,6 @@
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$SetUp
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 import org.junit.Before
 
 class A {<caret>

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpExists.kt.after

@@ -1,13 +1,9 @@
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$SetUp
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 import org.junit.Before
 
 class A {
-    @org.testng.annotations.BeforeMethod
-    fun setUp() {
-        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
-    }
-
     @Before
     fun setUp() {
         throw UnsupportedOperationException()

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpOverrides.kt

@@ -1,5 +1,6 @@
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$SetUp
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 open class A {
     open fun setUp() {
 

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpOverrides.kt.after

@@ -2,6 +2,7 @@ import org.junit.Before
 
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$SetUp
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 open class A {
     open fun setUp() {
 

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDown.kt

@@ -1,5 +1,6 @@
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$TearDown
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 class A {<caret>
 
 }

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDown.kt.after

@@ -2,6 +2,7 @@ import org.junit.After
 
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$TearDown
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 class A {
     @After
     fun tearDown() {

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDownExists.kt

@@ -1,5 +1,6 @@
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$TearDown
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 import org.junit.After
 
 class A {<caret>

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDownExists.kt.after

@@ -1,13 +1,9 @@
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$TearDown
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 import org.junit.After
 
 class A {
-    @org.testng.annotations.AfterMethod
-    fun tearDown() {
-        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
-    }
-
     @After
     fun tearDown() {
         throw UnsupportedOperationException()

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/testMethod.kt

@@ -1,5 +1,6 @@
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$Test
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 class A {<caret>
 
 }

idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/testMethod.kt.after

@@ -2,6 +2,7 @@ import org.junit.Test
 
 // ACTION_CLASS: org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase$Test
 // CONFIGURE_LIBRARY: JUnit@lib/junit-4.12.jar
+// TEST_FRAMEWORK: JUnit4
 class A {
     @Test
     fun name() {

idea/tests/org/jetbrains/kotlin/idea/codeInsight/generate/AbstractGenerateTestSupportMethodActionTest.kt

@@ -17,7 +17,9 @@
 package org.jetbrains.kotlin.idea.codeInsight.generate
 
 import com.intellij.openapi.roots.ModuleRootManager
+import org.jetbrains.kotlin.idea.actions.generate.KotlinGenerateTestSupportActionBase
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
+import org.jetbrains.kotlin.test.InTextDirectivesUtils
 
 abstract class AbstractGenerateTestSupportMethodActionTest : AbstractCodeInsightActionTest() {
     private fun setUpTestSourceRoot() {
@@ -33,6 +35,11 @@ abstract class AbstractGenerateTestSupportMethodActionTest : AbstractCodeInsight
         }
     }
 
+    override fun createAction(fileText: String) =
+            (super.createAction(fileText) as KotlinGenerateTestSupportActionBase).apply {
+                testFrameworkToUse = InTextDirectivesUtils.findStringWithPrefixes(fileText, "// TEST_FRAMEWORK:")
+            }
+
     override fun doTest(path: String) {
         setUpTestSourceRoot()
         super.doTest(path)