From 3bec910365f11796e0dff5f2d4ca6720a240fa3b Mon Sep 17 00:00:00 2001
From: Phodal Huang <h@phodal.com>
Date: Tue, 2 Jan 2024 16:38:16 +0800
Subject: [PATCH] feat(rust): try handle for Rust test

---
 gradle/libs.versions.toml                     |  2 +-
 .../unittest/rust/RustTestCodeService.kt      | 19 ++++++++++++++++++-
 .../unitmesh/pick/worker/lang/RustWorker.kt   |  1 +
 .../pick/SingleProjectCodePickerTest.kt       |  2 +-
 4 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 620dde1e..dba90ba9 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -17,7 +17,7 @@ clikt = "4.2.1"
 
 chocolate-factory = "0.4.5"
 
-chapi = "2.1.5"
+chapi = "2.2.0"
 archguard = "2.1.2"
 codedb = "0.1.3"
 
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/unittest/rust/RustTestCodeService.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/unittest/rust/RustTestCodeService.kt
index b1b617af..056e81be 100644
--- a/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/unittest/rust/RustTestCodeService.kt
+++ b/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/unittest/rust/RustTestCodeService.kt
@@ -13,7 +13,24 @@ class RustTestCodeService(val job: JobContext) : UnitTestService {
     override fun isApplicable(container: CodeContainer): Boolean = job.project.language == SupportedLang.RUST
 
     override fun build(container: CodeContainer): List<TypedTestIns> {
-        val testCode = container.DataStructures.filter { it.Module == "tests" }
+        val testDataStruct = container.DataStructures.filter {
+            it.Module == "tests"
+        }
+        if (testDataStruct.isEmpty()) return emptyList()
+
+        val testCode = testDataStruct.map { dataStruct ->
+            dataStruct.Functions.filter { function -> function.Annotations.any { it.Name == "test" } }
+        }.flatten()
+
+        if (testCode.isEmpty()) return emptyList()
+
+        println(testCode)
+        testCode.forEach {
+            it.FunctionCalls.forEach { functionCall ->
+                // todo: lookup method call in Chapi
+            }
+        }
+
         return emptyList()
     }
 
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/lang/RustWorker.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/lang/RustWorker.kt
index 25339e3f..2a0b926d 100644
--- a/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/lang/RustWorker.kt
+++ b/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/lang/RustWorker.kt
@@ -15,6 +15,7 @@ class RustWorker(override val context: WorkerContext) : LangWorker {
     override fun prepareJob(job: InstructionFileJob) {
         this.jobs.add(job)
         try {
+            println("RustWorker prepareJob: ${job.fileSummary.location}")
             val container = RustAnalyser().analysis(job.code, job.fileSummary.location)
             job.codeLines = job.code.lines()
             container.DataStructures.map { ds ->
diff --git a/unit-picker/src/test/kotlin/cc/unitmesh/pick/SingleProjectCodePickerTest.kt b/unit-picker/src/test/kotlin/cc/unitmesh/pick/SingleProjectCodePickerTest.kt
index 1d84eeb5..8136b658 100644
--- a/unit-picker/src/test/kotlin/cc/unitmesh/pick/SingleProjectCodePickerTest.kt
+++ b/unit-picker/src/test/kotlin/cc/unitmesh/pick/SingleProjectCodePickerTest.kt
@@ -60,7 +60,7 @@ class SingleProjectCodePickerTest {
         val picker = SingleProjectCodePicker(
             InsPickerOption(
                 language = "rust",
-                url = "https://github.com/feakin/fklang",
+                url = "https://github.com/unit-mesh/edge-infer",
                 maxTokenLength = 8192,
                 codeStrategyTypes = listOf(CodeStrategyType.RELATED_CODE),
                 completionTypes = listOf(