feat(provider): add imports to TestFileContext

Add the ability to include imports in the TestFileContext class. This allows for better tracking and management of imported packages in test files.
unit-mesh · Jan 10, 2024 · 05419e5 · 05419e5
1 parent 26714b2
commit 05419e5
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 6 deletions.
diff --git a/java/src/main/kotlin/cc/unitmesh/idea/service/JavaWriteTestService.kt b/java/src/main/kotlin/cc/unitmesh/idea/service/JavaWriteTestService.kt
@@ -7,6 +7,7 @@ import cc.unitmesh.devti.provider.WriteTestService
 import com.intellij.execution.configurations.RunProfile
 import com.intellij.lang.Language
 import com.intellij.openapi.application.ReadAction
+import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.fileEditor.FileDocumentManager
@@ -16,6 +17,7 @@ import com.intellij.openapi.vfs.LocalFileSystem
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileManager
 import com.intellij.psi.*
+import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.plugins.gradle.service.execution.GradleRunConfiguration
 import java.io.File
 
@@ -81,11 +83,16 @@ class JavaWriteTestService : WriteTestService() {
 
         project.guessProjectDir()?.refresh(true, true)
 
+        val imports = runReadAction {
+            val importList = PsiTreeUtil.getChildrenOfTypeAsList(sourceFile, PsiImportList::class.java)
+            importList.flatMap { it.allImportStatements.map { import -> import.text } }
+        }
+
         return if (testFile != null) {
-            TestFileContext(isNewFile, testFile, relatedModels, className, sourceFile.language, null)
+            TestFileContext(isNewFile, testFile, relatedModels, className, sourceFile.language, null, imports)
         } else {
             val targetFile = createTestFile(sourceFile, testDir!!, packageName, project)
-            TestFileContext(isNewFile = true, targetFile, relatedModels, "", sourceFile.language, null)
+            TestFileContext(isNewFile = true, targetFile, relatedModels, "", sourceFile.language, null, imports)
         }
     }
 

diff --git a/kotlin/src/main/kotlin/cc/unitmesh/kotlin/provider/KotlinWriteTestService.kt b/kotlin/src/main/kotlin/cc/unitmesh/kotlin/provider/KotlinWriteTestService.kt
@@ -89,12 +89,16 @@ class KotlinWriteTestService : WriteTestService() {
 
         project.guessProjectDir()?.refresh(true, true)
 
-        val currentClass =  runReadAction { ClassContextProvider(false).from(element) }
+        val currentClass = runReadAction { ClassContextProvider(false).from(element) }
+        val imports: List<String> = runReadAction {
+            (sourceFile as KtFile).importList?.imports?.map { it.importedFqName?.asString() ?: "" } ?: emptyList()
+        }
+
         return if (testFile != null) {
-            TestFileContext(isNewFile, testFile, relatedModels, className, sourceFile.language, currentClass)
+            TestFileContext(isNewFile, testFile, relatedModels, className, sourceFile.language, currentClass, imports)
         } else {
             val targetFile = createTestFile(sourceFile, testDir!!, packageName, project)
-            TestFileContext(isNewFile = true, targetFile, relatedModels, "", sourceFile.language, currentClass)
+            TestFileContext(isNewFile = true, targetFile, relatedModels, "", sourceFile.language, currentClass, imports)
         }
     }
 

diff --git a/src/main/kotlin/cc/unitmesh/devti/provider/WriteTestService.kt b/src/main/kotlin/cc/unitmesh/devti/provider/WriteTestService.kt
@@ -77,7 +77,6 @@ abstract class WriteTestService : LazyExtensionInstance<WriteTestService>() {
             val service = if (testServices.isNotEmpty()) {
                 testServices.first()
             } else {
-                // if lang == "TypeScript JSX", we just use TypeScript
                 val firstPartLang = lang.split(" ")[0]
                 val partLang = filterByLang(extensionList, firstPartLang)
                 if (partLang.isNotEmpty()) {

diff --git a/src/main/kotlin/cc/unitmesh/devti/provider/context/TestFileContext.kt b/src/main/kotlin/cc/unitmesh/devti/provider/context/TestFileContext.kt
@@ -11,4 +11,5 @@ data class TestFileContext(
     val testClassName: String?,
     val language: Language,
     val currentClass: ClassContext? = null,
+    val imports: List<String> = emptyList(),
 )