refactor: extract for test context

unit-mesh · Jul 28, 2023 · 65955cc · 65955cc
1 parent 71b485e
commit 65955cc
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 21 deletions.
diff --git a/java/src/main/kotlin/cc/unitmesh/idea/MvcUtil.kt b/java/src/main/kotlin/cc/unitmesh/idea/MvcUtil.kt
@@ -0,0 +1,16 @@
+package cc.unitmesh.idea
+
+object MvcUtil {
+    fun langSuffix(lang: String): String = when (lang.lowercase()) {
+        "java" -> "java"
+        "kotlin" -> "kt"
+        else -> "java"
+    }
+
+    fun isController(fileName: String, lang: String = "java"): Boolean {
+        return fileName.endsWith("Controller." + langSuffix(lang))
+    }
+
+    fun isService(fileName: String, lang: String = "java") =
+        fileName.endsWith("Service." + langSuffix(lang)) || fileName.endsWith("ServiceImpl." + langSuffix(lang))
+}
diff --git a/java/src/main/kotlin/cc/unitmesh/idea/flow/provider/JvmIdeaContextPrompter.kt b/java/src/main/kotlin/cc/unitmesh/idea/flow/provider/JvmIdeaContextPrompter.kt
@@ -9,6 +9,7 @@ import cc.unitmesh.devti.provider.context.ChatContextProvider
 import cc.unitmesh.devti.provider.context.ChatCreationContext
 import cc.unitmesh.devti.provider.context.ChatOrigin
 import cc.unitmesh.devti.settings.AutoDevSettingsState
+import cc.unitmesh.idea.MvcUtil
 import cc.unitmesh.idea.flow.MvcContextService
 import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.components.service
@@ -33,17 +34,6 @@ class JvmIdeaContextPrompter : ContextPrompter() {
         additionContext += context
     }
 
-    private fun langSuffix(): String = when (lang.lowercase()) {
-        "java" -> "java"
-        "kotlin" -> "kt"
-        else -> "java"
-    }
-
-    private fun isController() = fileName.endsWith("Controller." + langSuffix())
-    private fun isService() =
-        fileName.endsWith("Service." + langSuffix()) || fileName.endsWith("ServiceImpl." + langSuffix())
-
-
     override fun initContext(
         actionType: ChatActionType,
         selectedText: String,
@@ -132,15 +122,15 @@ class JvmIdeaContextPrompter : ContextPrompter() {
                 }
 
                 when {
-                    isController() -> {
+                    MvcUtil.isController(fileName, lang) -> {
                         val spec = CustomPromptConfig.load().spec["controller"]
                         if (!spec.isNullOrEmpty()) {
                             additionContext = "requirements: \n$spec"
                         }
                         additionContext += mvcContextService.controllerPrompt(file)
                     }
 
-                    isService() -> {
+                    MvcUtil.isService(fileName, lang) -> {
                         val spec = CustomPromptConfig.load().spec["service"]
                         if (!spec.isNullOrEmpty()) {
                             additionContext = "requirements: \n$spec"
@@ -213,7 +203,7 @@ class JvmIdeaContextPrompter : ContextPrompter() {
         val projectPath = project!!.basePath ?: ""
         runReadAction {
             val lookupFile = if (selectedText.contains(projectPath)) {
-                val regex = Regex("$projectPath(.*\\.)${langSuffix()}")
+                val regex = Regex("$projectPath(.*\\.)${MvcUtil.langSuffix(lang)}")
                 val relativePath = regex.find(selectedText)?.groupValues?.get(1) ?: ""
                 val file = LocalFileSystem.getInstance().findFileByPath(projectPath + relativePath)
                 file?.let { PsiManager.getInstance(project!!).findFile(it) }

diff --git a/java/src/main/kotlin/cc/unitmesh/idea/provider/TestContextProvider.kt b/java/src/main/kotlin/cc/unitmesh/idea/provider/TestContextProvider.kt
@@ -0,0 +1,53 @@
+package cc.unitmesh.idea.provider
+
+import cc.unitmesh.devti.gui.chat.ChatActionType
+import cc.unitmesh.devti.provider.context.ChatContextItem
+import cc.unitmesh.devti.provider.context.ChatContextProvider
+import cc.unitmesh.devti.provider.context.ChatCreationContext
+import cc.unitmesh.idea.MvcUtil
+import com.intellij.openapi.project.Project
+
+class TestContextProvider : ChatContextProvider {
+    override fun isApplicable(project: Project, creationContext: ChatCreationContext): Boolean {
+        return creationContext.action == ChatActionType.WRITE_TEST
+    }
+
+    override fun collect(project: Project, creationContext: ChatCreationContext): List<ChatContextItem> {
+        val items = mutableListOf<ChatContextItem>()
+
+        val isController = creationContext.sourceFile?.name?.let {
+            MvcUtil.isController(it)
+        } ?: false
+
+        val isService = creationContext.sourceFile?.name?.let {
+            MvcUtil.isService(it)
+        } ?: false
+
+        if (isController) {
+            items.add(
+                ChatContextItem(
+                    TestContextProvider::class,
+                    """You MUST use should_xx style for test method name.
+                    | You MUST use MockMvc and test API only.
+                    | You MUST use given-when-then style.
+                    | You MUST use should_xx style for test method name.""".trimMargin()
+                )
+            )
+        }
+
+        if (isService) {
+            items.add(
+                ChatContextItem(
+                    TestContextProvider::class,
+                    """You MUST use should_xx style for test method name.
+                    | You MUST use Mockito and test service only.
+                    | You MUST use given-when-then style.
+                    | You MUST use should_xx style for test method name. """.trimMargin()
+                )
+            )
+        }
+
+        return items
+    }
+
+}
diff --git a/java/src/main/resources/cc.unitmesh.idea.xml b/java/src/main/resources/cc.unitmesh.idea.xml
@@ -20,6 +20,7 @@
 
         <chatContextProvider implementation="cc.unitmesh.idea.provider.JavaVersionProvider"/>
         <chatContextProvider implementation="cc.unitmesh.idea.provider.SpringContextProvider"/>
+        <chatContextProvider implementation="cc.unitmesh.idea.provider.TestContextProvider"/>
 
         <contextPrompter
                 language="JAVA"

diff --git a/src/main/kotlin/cc/unitmesh/devti/intentions/task/TestCodeGenTask.kt b/src/main/kotlin/cc/unitmesh/devti/intentions/task/TestCodeGenTask.kt
@@ -55,16 +55,10 @@ class TestCodeGenTask(
         var prompter = if (testContext.isNewFile) {
             """Write unit test for following code. 
                                     | You MUST return code only, not explain.
-                                    | You MUST use given-when-then style.
-                                    | You MUST use should_xx style for test method name.
-                                    | When testing controller, you MUST use MockMvc and test API only.
                                     | """.trimMargin()
         } else {
             """Write unit test for following code. 
-                                    | You MUST return method code only, not java class, no explain.
-                                    | You MUST use given-when-then style.
-                                    | You MUST use should_xx style for test method name.
-                                    | When testing controller, you MUST use MockMvc and test API only.
+                                    | You MUST return method code only, no explain.
                                     | You MUST return start with @Test annotation.
                                     | """.trimMargin()
         }