feat: add for junit rules

unit-mesh · Dec 18, 2023 · dc5289d · dc5289d
1 parent e0084de
commit dc5289d
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 19 deletions.
diff --git a/java/src/main/kotlin/cc/unitmesh/idea/provider/JavaTestContextProvider.kt b/java/src/main/kotlin/cc/unitmesh/idea/provider/JavaTestContextProvider.kt
@@ -28,15 +28,15 @@ open class JavaTestContextProvider : ChatContextProvider {
 
         val isSpringRelated = creationContext.element?.let { isSpringRelated(it) } ?: false
 
-        val baseTestPrompt = """
+        var baseTestPrompt = """
             |You MUST use should_xx_xx style for test method name.
             |You MUST use given-when-then style.
             |- Test file should be complete and compilable, without need for further actions.
             |- Ensure that each test focuses on a single use case to maintain clarity and readability.
             |- Instead of using `@BeforeEach` methods for setup, include all necessary code initialization within each individual test method, do not write parameterized tests.
             |""".trimMargin()
 
-        // todo: check is spring project
+        baseTestPrompt += junitRule(project)
 
         items += when {
             isController && isSpringRelated -> {
@@ -64,6 +64,19 @@ open class JavaTestContextProvider : ChatContextProvider {
         return items
     }
 
+    private fun junitRule(project: Project): String {
+        SpringContextProvider.prepareLibraryData(project)?.forEach {
+            if (it.groupId?.contains("org.junit.jupiter") == true) {
+                return "| This project uses JUnit 5, you should import `org.junit.jupiter.api.Test` and use `@Test` annotation."
+            }
+            if (it.groupId?.contains("org.junit") == true) {
+                return "| This project uses JUnit 4, you should import `org.junit.Test` and use `@Test` annotation."
+            }
+        }
+
+        return ""
+    }
+
     private fun isSpringRelated(method: PsiElement): Boolean {
         when (method) {
             is PsiMethod -> {

diff --git a/java/src/main/kotlin/cc/unitmesh/idea/provider/SpringContextProvider.kt b/java/src/main/kotlin/cc/unitmesh/idea/provider/SpringContextProvider.kt
@@ -17,7 +17,7 @@ open class SpringContextProvider : ChatContextProvider {
     }
 
     private fun hasProjectLibraries(project: Project): Boolean {
-        prepareLibraryData(project)?.forEach {
+        Companion.prepareLibraryData(project)?.forEach {
             if (it.groupId?.contains("org.springframework") == true) {
                 return true
             }
@@ -68,23 +68,8 @@ open class SpringContextProvider : ChatContextProvider {
         )
     }
 
-    private fun prepareLibraryData(project: Project): List<LibraryData>? {
-        val basePath = project.basePath ?: return null
-        val projectData = ProjectDataManager.getInstance().getExternalProjectData(
-            project, GradleConstants.SYSTEM_ID, basePath
-        )
-
-        val libraryDataList = projectData?.externalProjectStructure?.children?.filter {
-            it.data is LibraryData
-        }?.map {
-            it.data as LibraryData
-        }
-
-        return libraryDataList
-    }
-
     private fun prepareLibrary(project: Project): TestStack {
-        val libraryDataList = prepareLibraryData(project)
+        val libraryDataList = Companion.prepareLibraryData(project)
 
         val testStack = TestStack()
         var hasMatchSpringMvc = false
@@ -139,4 +124,21 @@ open class SpringContextProvider : ChatContextProvider {
 
         return testStack
     }
+
+    companion object {
+        fun prepareLibraryData(project: Project): List<LibraryData>? {
+            val basePath = project.basePath ?: return null
+            val projectData = ProjectDataManager.getInstance().getExternalProjectData(
+                project, GradleConstants.SYSTEM_ID, basePath
+            )
+
+            val libraryDataList = projectData?.externalProjectStructure?.children?.filter {
+                it.data is LibraryData
+            }?.map {
+                it.data as LibraryData
+            }
+
+            return libraryDataList
+        }
+    }
 }