Skip to content

Commit 79aba96

Browse files
committed
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.
1 parent dd14143 commit 79aba96

File tree

4 files changed

+15
-3
lines changed

4 files changed

+15
-3
lines changed

java/src/test/kotlin/cc/unitmesh/idea/context/JavaClassContextTest.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ public class BlogService {
9797
"""'package: cc.unitmesh.untitled.demo.controller.BlogController
9898
'@Controller
9999
class BlogController {
100-
blogService
100+
'variable -> BlogService blogService;
101101
+ public BlogController(BlogService blogService)
102102
+ @PostMapping("/blog") public BlogPost createBlog(CreateBlogDto blogDto)
103103
+ @GetMapping("/blog") public List<BlogPost> getBlog()

rust/src/main/kotlin/cc/unitmesh/rust/provider/RustTestService.kt

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,8 @@ class RustTestService : WriteTestService() {
7878
elementName,
7979
RsLanguage,
8080
currentObject,
81-
imports
81+
imports,
82+
element
8283
)
8384
}
8485

src/main/kotlin/cc/unitmesh/devti/intentions/action/task/TestCodeGenTask.kt

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import com.intellij.openapi.progress.Task
2222
import com.intellij.openapi.project.DumbService
2323
import com.intellij.openapi.project.Project
2424
import com.intellij.openapi.vfs.VirtualFile
25+
import com.intellij.psi.PsiNameIdentifierOwner
2526
import kotlinx.coroutines.flow.*
2627
import kotlinx.coroutines.runBlocking
2728

@@ -102,7 +103,12 @@ class TestCodeGenTask(val request: TestCodeGenRequest) :
102103
"$comment $it"
103104
}
104105

105-
testPromptContext.sourceCode = request.selectText
106+
testPromptContext.sourceCode = if(request.element !is PsiNameIdentifierOwner) {
107+
testContext.testElement?.text ?: ""
108+
} else {
109+
request.selectText
110+
}
111+
106112
testPromptContext.isNewFile = testContext.isNewFile
107113

108114
templateRender.context = testPromptContext

src/main/kotlin/cc/unitmesh/devti/provider/context/TestFileContext.kt

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package cc.unitmesh.devti.provider.context
33
import cc.unitmesh.devti.context.ClassContext
44
import com.intellij.lang.Language
55
import com.intellij.openapi.vfs.VirtualFile
6+
import com.intellij.psi.PsiElement
67

78
data class TestFileContext(
89
val isNewFile: Boolean,
@@ -17,4 +18,8 @@ data class TestFileContext(
1718
*/
1819
val currentObject: String? = null,
1920
val imports: List<String> = emptyList(),
21+
/** Since 1.5.4, since some languages have different test code insertion strategies,
22+
* we need to pass in the test element text
23+
*/
24+
val testElement: PsiElement? = null,
2025
)

0 commit comments

Comments
 (0)