feat: add test support for mod

phodal · Jan 1, 2024 · da588ce · da588ce
1 parent 06333ed
commit da588ce
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 0 deletions.
diff --git a/chapi-ast-rust/src/main/kotlin/chapi/ast/rustast/RustAstBaseListener.kt b/chapi-ast-rust/src/main/kotlin/chapi/ast/rustast/RustAstBaseListener.kt
@@ -29,6 +29,10 @@ open class RustAstBaseListener(private val fileName: String) : RustParserBaseLis
 
     private var structMap = mutableMapOf<String, CodeDataStruct>()
 
+    private var currentModule: String = ""
+    /// for testing
+    private var lastModule: String = ""
+
 
     /**
      * localVars will store all local variables in the current scope
@@ -82,6 +86,15 @@ open class RustAstBaseListener(private val fileName: String) : RustParserBaseLis
         Imports = imports
     }
 
+    override fun enterModule(ctx: RustParser.ModuleContext?) {
+        currentModule = (ctx?.identifier()?.text ?: return)
+    }
+
+    override fun exitModule(ctx: RustParser.ModuleContext?) {
+        lastModule = currentModule
+        currentModule = ""
+    }
+
     override fun enterUseDeclaration(ctx: RustParser.UseDeclarationContext?) {
         val importString: List<List<String>> = buildToPath(ctx)
 
@@ -150,6 +163,7 @@ open class RustAstBaseListener(private val fileName: String) : RustParserBaseLis
 
         val codeStruct = CodeDataStruct(
             NodeName = structName,
+            Module = currentModule,
             Package = codeContainer.PackageName,
             Annotations = annotation,
             Fields = buildFields(ctx.structFields()),
@@ -270,6 +284,7 @@ open class RustAstBaseListener(private val fileName: String) : RustParserBaseLis
 
         val codeStruct = CodeDataStruct(
             NodeName = enumName,
+            Module = currentModule,
             Package = codeContainer.PackageName,
             Annotations = annotation,
             Fields = buildEnumFields(ctx.enumItems()),
@@ -367,6 +382,7 @@ open class RustAstBaseListener(private val fileName: String) : RustParserBaseLis
         } else {
             currentNode = CodeDataStruct(
                 NodeName = nodeName,
+                Module = currentModule,
                 Package = codeContainer.PackageName,
                 Position = buildPosition(ctx ?: return)
             )
@@ -417,6 +433,7 @@ open class RustAstBaseListener(private val fileName: String) : RustParserBaseLis
         return listOf(
             CodeDataStruct().apply {
                 NodeName = fileName.substringBeforeLast('.')
+                Module = if (lastModule == "tests") lastModule else ""
                 Type = DataStructType.OBJECT
                 Package = codeContainer.PackageName
                 FilePath = codeContainer.FullName

diff --git a/chapi-ast-rust/src/test/kotlin/chapi/ast/rustast/RustFullIdentListenerTest.kt b/chapi-ast-rust/src/test/kotlin/chapi/ast/rustast/RustFullIdentListenerTest.kt
@@ -421,4 +421,23 @@ class RustFullIdentListenerTest {
         assertEquals("actix_web::HttpServer", secondFunction.FunctionCalls[0].NodeName)
         assertEquals("run", secondFunction.FunctionCalls[0].FunctionName)
     }
+
+    @Test
+    fun should_handle_test_mod() {
+        val code = """
+            #[cfg(test)]
+            mod tests {
+                use super::*;
+            
+                #[test]
+                fn test_add() {
+                    assert_eq!(add(1, 2), 3);
+                }
+            }
+        """.trimIndent()
+
+        val codeContainer = RustAnalyser().analysis(code, "lib.rs")
+        val codeDataStruct = codeContainer.DataStructures[0]
+        assertEquals("tests", codeDataStruct.Module)
+    }
 }
diff --git a/chapi-domain/src/main/kotlin/chapi/domain/core/CodeDataStruct.kt b/chapi-domain/src/main/kotlin/chapi/domain/core/CodeDataStruct.kt
@@ -27,6 +27,7 @@ data class CodeDataStruct(
      * for Rust, if is a variable, function, it will be naming to `default`
      */
     var NodeName: String = "",
+    //
     var Module: String = "",
     var Type: DataStructType = DataStructType.EMPTY,
     var Package: String = "",