feat(provider): add test element to TestFileContext

Add the `testElement` property to the `TestFileContext` data class in order to pass in the test element text. This is necessary since some languages have different test code insertion strategies. The `testElement` property is of type `PsiElement` and is set to `null` by default.
unit-mesh · Jan 19, 2024 · 79aba96 · 79aba96
1 parent dd14143
commit 79aba96
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 3 deletions.
diff --git a/java/src/test/kotlin/cc/unitmesh/idea/context/JavaClassContextTest.kt b/java/src/test/kotlin/cc/unitmesh/idea/context/JavaClassContextTest.kt
@@ -97,7 +97,7 @@ public class BlogService {
             """'package: cc.unitmesh.untitled.demo.controller.BlogController
 '@Controller
 class BlogController {
-  blogService
+  'variable -> BlogService blogService;
   + public BlogController(BlogService blogService)
   + @PostMapping("/blog")     public BlogPost createBlog(CreateBlogDto blogDto)
   + @GetMapping("/blog")     public List<BlogPost> getBlog()

diff --git a/rust/src/main/kotlin/cc/unitmesh/rust/provider/RustTestService.kt b/rust/src/main/kotlin/cc/unitmesh/rust/provider/RustTestService.kt
@@ -78,7 +78,8 @@ class RustTestService : WriteTestService() {
             elementName,
             RsLanguage,
             currentObject,
-            imports
+            imports,
+            element
         )
     }
 

diff --git a/src/main/kotlin/cc/unitmesh/devti/intentions/action/task/TestCodeGenTask.kt b/src/main/kotlin/cc/unitmesh/devti/intentions/action/task/TestCodeGenTask.kt
@@ -22,6 +22,7 @@ import com.intellij.openapi.progress.Task
 import com.intellij.openapi.project.DumbService
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiNameIdentifierOwner
 import kotlinx.coroutines.flow.*
 import kotlinx.coroutines.runBlocking
 
@@ -102,7 +103,12 @@ class TestCodeGenTask(val request: TestCodeGenRequest) :
             "$comment $it"
         }
 
-        testPromptContext.sourceCode = request.selectText
+        testPromptContext.sourceCode = if(request.element !is PsiNameIdentifierOwner) {
+            testContext.testElement?.text ?: ""
+        } else {
+            request.selectText
+        }
+
         testPromptContext.isNewFile = testContext.isNewFile
 
         templateRender.context = testPromptContext

diff --git a/src/main/kotlin/cc/unitmesh/devti/provider/context/TestFileContext.kt b/src/main/kotlin/cc/unitmesh/devti/provider/context/TestFileContext.kt
@@ -3,6 +3,7 @@ package cc.unitmesh.devti.provider.context
 import cc.unitmesh.devti.context.ClassContext
 import com.intellij.lang.Language
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiElement
 
 data class TestFileContext(
     val isNewFile: Boolean,
@@ -17,4 +18,8 @@ data class TestFileContext(
      */
     val currentObject: String? = null,
     val imports: List<String> = emptyList(),
+    /** Since 1.5.4, since some languages have different test code insertion strategies,
+     * we need to pass in the test element text
+     */
+    val testElement: PsiElement? = null,
 )