feat(kotlin): init test services

unit-mesh · Dec 28, 2023 · 5d26c16 · 5d26c16
1 parent e03dd14
commit 5d26c16
Show file tree

Hide file tree

Showing 7 changed files with 54 additions and 11 deletions.
diff --git a/unit-core/src/main/kotlin/cc/unitmesh/core/SupportedLang.kt b/unit-core/src/main/kotlin/cc/unitmesh/core/SupportedLang.kt
@@ -3,19 +3,21 @@ package cc.unitmesh.core
 enum class SupportedLang(val extension: String) {
     JAVA("java"),
     TYPESCRIPT("ts"),
+    KOTLIN("kt"),
     ;
 
     companion object {
         fun from(value: String): SupportedLang? {
             return when (value.lowercase()) {
                 "java" -> JAVA
                 "typescript" -> TYPESCRIPT
+                "kotlin" -> KOTLIN
                 else -> null
             }
         }
 
         fun all(): List<SupportedLang> {
-            return listOf(JAVA, TYPESCRIPT)
+            return listOf(JAVA, TYPESCRIPT, KOTLIN)
         }
     }
 }
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/SingleProjectCodePicker.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/SingleProjectCodePicker.kt
@@ -3,6 +3,7 @@ package cc.unitmesh.pick
 import cc.unitmesh.pick.worker.job.InstructionFileJob
 import cc.unitmesh.pick.option.InsPickerOption
 import cc.unitmesh.core.Instruction
+import cc.unitmesh.core.SupportedLang
 import cc.unitmesh.pick.ext.GitUtil
 import cc.unitmesh.pick.ext.PickDirectoryWalker
 import cc.unitmesh.pick.threshold.InsQualityThreshold
@@ -84,7 +85,10 @@ class SingleProjectCodePicker(private val config: InsPickerOption) {
                 )
             )
         )
-        workerManager.init(codeDir, config.language)
+
+
+        val language = SupportedLang.from(config.language.lowercase()) ?: throw IllegalArgumentException("unsupported language: ${config.language}")
+        workerManager.init(codeDir, language)
 
         val walkdirChannel = Channel<FileJob>()
         val summary = mutableListOf<Instruction>()
@@ -97,7 +101,7 @@ class SingleProjectCodePicker(private val config: InsPickerOption) {
             launch {
                 for (fileJob in walkdirChannel) {
                     languageWorker.processFile(fileJob)?.let {
-                        workerManager.tryAddJobByThreshold(InstructionFileJob.from(it))
+                        workerManager.tryAddJobByThreshold(InstructionFileJob.from(it), language)
                     }
                 }
 

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/unittest/lang/JavaTestCodeService.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/unittest/lang/JavaTestCodeService.kt
@@ -11,7 +11,7 @@ import chapi.domain.core.CodeDataStruct
 import chapi.domain.core.CodeFunction
 import chapi.domain.core.CodeImport
 
-class JavaTestCodeService(val context: JobContext) : UnitTestService {
+open class JavaTestCodeService(open val context: JobContext) : UnitTestService {
     override fun isApplicable(dataStruct: CodeDataStruct): Boolean {
         return dataStruct.NodeName.endsWith("Test") || dataStruct.NodeName.endsWith("Tests")
     }

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/unittest/lang/KotlinTestCodeService.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/unittest/lang/KotlinTestCodeService.kt
@@ -0,0 +1,6 @@
+package cc.unitmesh.pick.builder.unittest.lang
+
+import cc.unitmesh.pick.worker.job.JobContext
+
+class KotlinTestCodeService(override val context: JobContext) : JavaTestCodeService(context) {
+}
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/unittest/lang/UnitTestService.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/unittest/lang/UnitTestService.kt
@@ -31,6 +31,7 @@ interface UnitTestService {
                 when (it) {
                     SupportedLang.JAVA -> JavaTestCodeService(job)
                     SupportedLang.TYPESCRIPT -> TypeScriptTestCodeService(job)
+                    SupportedLang.KOTLIN -> KotlinTestCodeService(job)
                 }
             }
 

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/WorkerManager.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/WorkerManager.kt
@@ -20,7 +20,8 @@ import java.util.*
 class WorkerManager(private val context: WorkerContext) {
     private val workers: Map<SupportedLang, LangWorker> = mapOf(
         SupportedLang.JAVA to JavaWorker(context),
-//        SupportedLang.TYPESCRIPT to TypescriptWorker(context),
+        SupportedLang.KOTLIN to JavaWorker(context),
+        SupportedLang.TYPESCRIPT to TypescriptWorker(context),
     )
 
     private val thresholdChecker: ThresholdChecker = ThresholdChecker(context)
@@ -34,11 +35,11 @@ class WorkerManager(private val context: WorkerContext) {
      * @param language The programming language used in the project.
      *
      */
-    fun init(codeDir: File, language: String) {
+    fun init(codeDir: File, language: SupportedLang) {
         try {
             val dependencies = ScaAnalyser(object : ScaContext {
                 override val client: ArchGuardClient = EmptyArchGuardClient()
-                override val language: String = language
+                override val language: String = language.name.lowercase()
                 override val path: String = codeDir.absolutePath
             }).analyse()
 
@@ -61,13 +62,17 @@ class WorkerManager(private val context: WorkerContext) {
      *
      * @return true if the job is successfully added to the worker, false otherwise.
      */
-    fun tryAddJobByThreshold(job: InstructionFileJob): Boolean {
+    fun tryAddJobByThreshold(job: InstructionFileJob, supportedLangs: SupportedLang): Boolean {
         if (!thresholdChecker.isMetThreshold(job)) {
             return false
         }
 
-        val language = SupportedLang.from(job.fileSummary.language)
-        val worker = workers[language] ?: return false
+        val parsedLanguage = SupportedLang.from(job.fileSummary.language) ?: return false
+        if (parsedLanguage != supportedLangs) {
+            return false
+        }
+
+        val worker = workers[parsedLanguage] ?: return false
         worker.prepareJob(job)
 
         return true

diff --git a/unit-picker/src/test/kotlin/cc/unitmesh/pick/SingleProjectCodePickerTest.kt b/unit-picker/src/test/kotlin/cc/unitmesh/pick/SingleProjectCodePickerTest.kt
@@ -1,6 +1,7 @@
 package cc.unitmesh.pick;
 
 import cc.unitmesh.core.Instruction
+import cc.unitmesh.core.completion.CompletionBuilderType
 import cc.unitmesh.pick.option.InsPickerOption
 import kotlinx.coroutines.runBlocking
 import kotlinx.serialization.encodeToString
@@ -15,7 +16,7 @@ class SingleProjectCodePickerTest {
         val picker = SingleProjectCodePicker(
             InsPickerOption(
                 url = "https://github.com/unit-mesh/unit-eval-testing",
-                completionTypeSize  = 10,
+                completionTypeSize = 10,
                 maxCharInCode = 100,
             )
         )
@@ -28,4 +29,28 @@ class SingleProjectCodePickerTest {
             })
         }
     }
+
+    @Test
+    fun shouldCheckoutTestCodeWithBadSmell() {
+        val picker = SingleProjectCodePicker(
+            InsPickerOption(
+                language = "kotlin",
+                url = "https://github.com/unit-mesh/unit-eval",
+                completionTypeSize = 10,
+                maxCharInCode = 100,
+                completionTypes = listOf(
+                    CompletionBuilderType.IN_BLOCK_COMPLETION,
+//                    CompletionBuilderType.TEST_CODE_GEN, CompletionBuilderType.DOCUMENTATION
+                ),
+            )
+        )
+
+        val outputFile = File("test.jsonl")
+        runBlocking {
+            val output: MutableList<Instruction> = picker.execute()
+            outputFile.writeText(output.joinToString("\n") {
+                Json.encodeToString(it)
+            })
+        }
+    }
 }