From 1e37eaa56b4a2a3518d2ed00621a346e588e0f70 Mon Sep 17 00:00:00 2001
From: zhislin <zhislin@icloud.com>
Date: Fri, 15 Nov 2024 13:29:03 +0800
Subject: [PATCH] fix the comment info collection

---
 .../CheckAnnotationCommand.kt                 | 20 +++++++++++++++++
 .../AnnotationCompletionContributor.kt        | 22 ++++++++++++-------
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/src/main/kotlin/com/github/jaksonlin/pitestintellij/commands/unittestannotations/CheckAnnotationCommand.kt b/src/main/kotlin/com/github/jaksonlin/pitestintellij/commands/unittestannotations/CheckAnnotationCommand.kt
index e547257..2302b14 100644
--- a/src/main/kotlin/com/github/jaksonlin/pitestintellij/commands/unittestannotations/CheckAnnotationCommand.kt
+++ b/src/main/kotlin/com/github/jaksonlin/pitestintellij/commands/unittestannotations/CheckAnnotationCommand.kt
@@ -6,6 +6,9 @@ import com.github.jaksonlin.pitestintellij.context.CaseCheckContext
 import com.github.jaksonlin.pitestintellij.context.UnittestCase
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiAnnotation
+import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiMethod
+import com.intellij.psi.util.PsiTreeUtil
 
 class CheckAnnotationCommand(project: Project,  context: CaseCheckContext):UnittestCaseCheckCommand(project, context) {
     override fun execute() {
@@ -32,6 +35,22 @@ class CheckAnnotationCommand(project: Project,  context: CaseCheckContext):Unitt
         }
     }
 
+    private fun extractMethodBodyComments(psiMethod: PsiMethod): String {
+        val stepComments = mutableListOf<String>()
+        val assertComments = mutableListOf<String>()
+
+        PsiTreeUtil.findChildrenOfType(psiMethod, PsiComment::class.java).forEach {
+            if (it.text.contains("step", ignoreCase = true)) {
+                stepComments.add(it.text)
+            }
+            if (it.text.contains("assert", ignoreCase = true)) {
+                assertComments.add(it.text)
+            }
+        }
+        // format as : test_step_1: step_comment, remove the leading //
+        return stepComments.mapIndexed { index, comment -> "test_step_${index + 1}: ${comment.substring(2)}" }.joinToString("\n") + "\n" + assertComments.mapIndexed { index, comment -> "test_assert_${index + 1}: ${comment.substring(2)}" }.joinToString("\n")
+    }
+
 
     private fun formatTestCaseMessage(
         testCase: UnittestCase,
@@ -48,6 +67,7 @@ class CheckAnnotationCommand(project: Project,  context: CaseCheckContext):Unitt
                     appendLine(testCase.getStringList(field.name).joinToString(", "))
             }
         }
+        appendLine(extractMethodBodyComments(context.psiMethod))
     }
 
 }
\ No newline at end of file
diff --git a/src/main/kotlin/com/github/jaksonlin/pitestintellij/completion/AnnotationCompletionContributor.kt b/src/main/kotlin/com/github/jaksonlin/pitestintellij/completion/AnnotationCompletionContributor.kt
index 48ee9f7..6426bf8 100644
--- a/src/main/kotlin/com/github/jaksonlin/pitestintellij/completion/AnnotationCompletionContributor.kt
+++ b/src/main/kotlin/com/github/jaksonlin/pitestintellij/completion/AnnotationCompletionContributor.kt
@@ -58,8 +58,8 @@ class AnnotationCompletionContributor : CompletionContributor() {
                     LOG.info("Actual annotation class: ${annotation.qualifiedName}")
 
                     // Check if this is our target annotation
-                    if (annotation.qualifiedName != schema.annotationClassName) {
-                        LOG.info("Annotation mismatch")
+                    if (annotation.qualifiedName?.endsWith(schema.annotationClassName) != true) {
+                        LOG.info("Annotation mismatch: ${annotation.qualifiedName} not the same as ${schema.annotationClassName}")
                         return
                     }
 
@@ -86,13 +86,19 @@ class AnnotationCompletionContributor : CompletionContributor() {
                             fieldType = field.type,
                             isDefaultValue = isDefault
                         )
-
-                        val prioritized = when {
-                            isDefault -> PrioritizedLookupElement.withPriority(element, 100.0)
-                            field.validation.mode == ValidationMode.EXACT ->
-                                PrioritizedLookupElement.withPriority(element, 50.0)
-                            else -> PrioritizedLookupElement.withPriority(element, 0.0)
+                        var properitizedValue = 100.0
+                        when {
+                            // Highest priority for exact matches
+                            result.prefixMatcher.prefixMatches(value) -> properitizedValue = 100.0
+                            // High priority for default values
+                            isDefault -> properitizedValue = 90.0
+                            // Medium priority for validated values
+                            field.validation?.validValues?.contains(value) == true -> properitizedValue = 80.0
+                            // Lower priority for other suggestions
+                            else -> properitizedValue = 70.0
                         }
+                        val prioritized = PrioritizedLookupElement.withPriority(element,properitizedValue)
+                        LOG.info("Adding element: $value with priority ${properitizedValue}")
 
                         result.addElement(prioritized)
                     }