diff --git a/idea/src/org/jetbrains/kotlin/idea/actions/generate/KotlinGenerateTestSupportActionBase.kt b/idea/src/org/jetbrains/kotlin/idea/actions/generate/KotlinGenerateTestSupportActionBase.kt
index dac27e3e21b..9316fc5daf8 100644
--- a/idea/src/org/jetbrains/kotlin/idea/actions/generate/KotlinGenerateTestSupportActionBase.kt
+++ b/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) }
-        chooseAndPerform(editor, frameworks) { doGenerate(editor, file, klass, it) }
+
+        if (testFrameworkToUse != null) {
+            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) {
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/dataMethod.kt b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/dataMethod.kt
index 1fec155600d..a2a46034c91 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/dataMethod.kt
+++ b/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>
 
 }
\ No newline at end of file
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/dataMethod.kt.after b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/dataMethod.kt.after
index a3cea12dfd7..8948aad2a8d 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/dataMethod.kt.after
+++ b/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>> {
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUp.kt b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUp.kt
index cdb241b0f27..c02aeac3d64 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUp.kt
+++ b/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>
 
 }
\ No newline at end of file
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUp.kt.after b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUp.kt.after
index 663bc141a3f..8d3b6f6fdeb 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUp.kt.after
+++ b/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() {
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpExists.kt b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpExists.kt
index c260bde9f3a..bd599cead5b 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpExists.kt
+++ b/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>
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpExists.kt.after b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpExists.kt.after
index a30032869ee..380a938c435 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpExists.kt.after
+++ b/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()
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpOverrides.kt b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpOverrides.kt
index 5db6146900c..d0c69ec5120 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpOverrides.kt
+++ b/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() {
 
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpOverrides.kt.after b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpOverrides.kt.after
index 2eaf72052e4..11e5a5a5d56 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/setUpOverrides.kt.after
+++ b/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() {
 
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDown.kt b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDown.kt
index c9ac2501bd7..d41f66520f9 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDown.kt
+++ b/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>
 
 }
\ No newline at end of file
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDown.kt.after b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDown.kt.after
index b3d7f5a84bd..9924210c27b 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDown.kt.after
+++ b/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() {
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDownExists.kt b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDownExists.kt
index ac5a7798595..ceaf0f20bab 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDownExists.kt
+++ b/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>
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDownExists.kt.after b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDownExists.kt.after
index da36100a7d2..7ca52c81388 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/tearDownExists.kt.after
+++ b/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()
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/testMethod.kt b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/testMethod.kt
index 6d63f9cb750..51c4fb3435d 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/testMethod.kt
+++ b/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>
 
 }
\ No newline at end of file
diff --git a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/testMethod.kt.after b/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/testMethod.kt.after
index 3c08dfe45a0..801ccc3f094 100644
--- a/idea/testData/codeInsight/generate/testFrameworkSupport/jUnit4/testMethod.kt.after
+++ b/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() {
diff --git a/idea/tests/org/jetbrains/kotlin/idea/codeInsight/generate/AbstractGenerateTestSupportMethodActionTest.kt b/idea/tests/org/jetbrains/kotlin/idea/codeInsight/generate/AbstractGenerateTestSupportMethodActionTest.kt
index c41a2211c7d..5f408daa907 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/codeInsight/generate/AbstractGenerateTestSupportMethodActionTest.kt
+++ b/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)