Test Framework Support: Allow using whitespaces and other symbols in "Generate -> Test Function" dialog

#KT-12556

ChangeLog.md

@@ -196,6 +196,7 @@ These artifacts include extensions for the types available in the latter JDKs, s
 - [`KT-13838`](https://youtrack.jetbrains.com/issue/KT-13838) Add file name to the presentation of private top-level declaration (Go to symbol, etc.)
 - [`KT-14096`](https://youtrack.jetbrains.com/issue/KT-14096) Rename: When renaming Kotlin file outside of source root do not rename its namesake in a source root
 - [`KT-13928`](https://youtrack.jetbrains.com/issue/KT-13928) Move Inner Class to Upper Level: Fix replacement of outer class instances used in inner class constructor calls
+- [`KT-12556`](https://youtrack.jetbrains.com/issue/KT-12556) Allow using whitespaces and other symbols in "Generate -> Test Function" dialog
 
 #### Intention actions, inspections and quickfixes
 

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

@@ -43,17 +43,16 @@ import org.jetbrains.kotlin.asJava.toLightClass
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.idea.core.KotlinNameSuggester
+import org.jetbrains.kotlin.idea.core.insertMember
 import org.jetbrains.kotlin.idea.core.overrideImplement.OverrideMemberChooserObject.BodyType
 import org.jetbrains.kotlin.idea.core.overrideImplement.generateUnsupportedOrSuperCall
-import org.jetbrains.kotlin.idea.refactoring.j2k
+import org.jetbrains.kotlin.idea.core.quoteIfNeeded
 import org.jetbrains.kotlin.idea.quickfix.createFromUsage.callableBuilder.setupEditorSelection
-import org.jetbrains.kotlin.idea.core.insertMember
+import org.jetbrains.kotlin.idea.refactoring.j2k
 import org.jetbrains.kotlin.idea.testIntegration.findSuitableFrameworks
 import org.jetbrains.kotlin.idea.util.application.executeWriteCommand
 import org.jetbrains.kotlin.lexer.KtTokens
-import org.jetbrains.kotlin.psi.KtClassOrObject
-import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.kotlin.utils.ifEmpty
 
@@ -88,7 +87,7 @@ abstract class KotlinGenerateTestSupportActionBase(
         private val NAME_VAR = "\${NAME}"
 
         private val NAME_VALIDATOR = object : InputValidator {
-            override fun checkInput(inputString: String) = KotlinNameSuggester.isIdentifier(inputString)
+            override fun checkInput(inputString: String) = KotlinNameSuggester.isIdentifier(inputString.quoteIfNeeded())
             override fun canClose(inputString: String) = true
         }
     }
@@ -133,6 +132,8 @@ abstract class KotlinGenerateTestSupportActionBase(
         chooseAndPerform(editor, frameworks) { doGenerate(editor, file, klass, it) }
     }
 
+    private val DUMMY_NAME = "__KOTLIN_RULEZZZ__"
+
     private fun doGenerate(editor: Editor, file: PsiFile, klass: KtClassOrObject, framework: TestFramework) {
         val project = file.project
         val commandName = "Generate test function"
@@ -143,23 +144,27 @@ abstract class KotlinGenerateTestSupportActionBase(
                 val fileTemplateDescriptor = methodKind.getFileTemplateDescriptor(framework)
                 val fileTemplate = FileTemplateManager.getInstance(project).getCodeTemplate(fileTemplateDescriptor.fileName)
                 var templateText = fileTemplate.text.replace(BODY_VAR, "")
+                var name: String? = null
                 if (templateText.contains(NAME_VAR)) {
-                    var name = if (templateText.contains("test$NAME_VAR")) "Name" else "name"
+                    name = if (templateText.contains("test$NAME_VAR")) "Name" else "name"
                     if (!ApplicationManager.getApplication().isUnitTestMode) {
                         name = Messages.showInputDialog("Choose test name: ", commandName, null, name, NAME_VALIDATOR)
                                 ?: return@executeWriteCommand
                     }
 
-                    templateText = fileTemplate.text.replace(NAME_VAR, name)
+                    templateText = fileTemplate.text.replace(NAME_VAR, DUMMY_NAME)
                 }
                 val factory = PsiElementFactory.SERVICE.getInstance(project)
                 val psiMethod = factory.createMethodFromText(templateText, null)
                 psiMethod.throwsList.referenceElements.forEach { it.delete() }
-                val function = psiMethod.j2k() as? KtNamedFunction
+                var function = psiMethod.j2k() as? KtNamedFunction
                 if (function == null) {
                     HintManager.getInstance().showErrorHint(editor, "Couldn't convert Java template to Kotlin")
                     return@executeWriteCommand
                 }
+                if (name != null) {
+                    function = substituteNewName(function, name)
+                }
                 val functionInPlace = insertMember(editor, klass, function)
 
                 val functionDescriptor = functionInPlace.resolveToDescriptor() as FunctionDescriptor
@@ -184,6 +189,38 @@ abstract class KotlinGenerateTestSupportActionBase(
         }
     }
 
+    private fun substituteNewName(function: KtNamedFunction, name: String): KtNamedFunction {
+        val psiFactory = KtPsiFactory(function)
+
+        // First replace all DUMMY_NAME occurrences in names as they need special treatment due to quotation
+        var function1 = function
+        function1.accept(
+                object : KtTreeVisitorVoid() {
+                    private fun getNewId(currentId: String): String? {
+                        if (!currentId.contains(DUMMY_NAME)) return null
+                        return currentId.replace(DUMMY_NAME, name).quoteIfNeeded()
+                    }
+
+                    override fun visitNamedDeclaration(declaration: KtNamedDeclaration) {
+                        val nameIdentifier = declaration.nameIdentifier ?: return
+                        val newId = getNewId(nameIdentifier.text) ?: return
+                        declaration.setName(newId)
+                    }
+
+                    override fun visitSimpleNameExpression(expression: KtSimpleNameExpression) {
+                        val newId = getNewId(expression.text) ?: return
+                        expression.replace(psiFactory.createSimpleName(newId))
+                    }
+                }
+        )
+        // Then text-replace remaining occurrences (if any)
+        val functionText = function1.text
+        if (functionText.contains(DUMMY_NAME)) {
+            function1 = psiFactory.createFunction(function1.text.replace(DUMMY_NAME, name))
+        }
+        return function1
+    }
+
     override fun createEditTemplateAction(dataContext: DataContext): AnAction? {
         val project = CommonDataKeys.PROJECT.getData(dataContext) ?: return null
         val editor = CommonDataKeys.EDITOR.getData(dataContext) ?: return null