some test for import classifier

Author: deadlovelll

src/core/ast_parser/ast_parser.rs

@@ -7,14 +7,12 @@ use std::io;
 use tree_sitter::{Language};
 use tree_sitter_python;
 
-use crate::core::ast_parser::import_classifier::ImportClassifier;
 use crate::core::ast_parser::query::ParseQuery;
 use crate::core::ast_parser::syntax_tree::SyntaxTree;
 use crate::core::ast_parser::tree_analyzer::TreeAnalyzer;
 
 pub struct AstParser {
     file_data: Mutex<HashMap<String, HashMap<String, Vec<String>>>>,
-    import_classifier: ImportClassifier,
     query: ParseQuery,
     tree: SyntaxTree,
     tree_analyzer: TreeAnalyzer,
@@ -25,7 +23,6 @@ impl AstParser {
     pub fn new() -> Self {
         Self {
             file_data: Mutex::new(HashMap::new()),
-            import_classifier: ImportClassifier::new(),
             query: ParseQuery::new(),
             tree: SyntaxTree::new(),
             tree_analyzer: TreeAnalyzer::new(),

src/core/ast_parser/import_classifier.rs

@@ -29,10 +29,57 @@ impl ImportClassifier {
     ) -> bool {
 
         for dir in root_dirs {
-            if import.starts_with(dir) {
+            if import.starts_with(dir) || import.starts_with(".") {
                 return true;
             }
         }
         return false;
     }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_import_is_local() {
+        let root_dirs = vec![
+            "application".to_string(), 
+            "services".to_string()
+        ];
+        let import = "services.foo.Bar".to_string();
+
+        let classifier = ImportClassifier::new();
+        let result = classifier.is_eligible(&import, &root_dirs);
+
+        assert_eq!(result, true);
+    }
+
+    #[test]
+    fn test_import_is_not_local() {
+        let root_dirs = vec![
+            "application".to_string(), 
+            "services".to_string()
+        ];
+        let import = "fastapi.ApiRouter".to_string();
+
+        let classifier = ImportClassifier::new();
+        let result = classifier.is_eligible(&import, &root_dirs);
+
+        assert_eq!(result, false);
+    }
+
+    #[test]
+    fn test_import_is_relative() {
+        let root_dirs = vec![
+            "application".to_string(), 
+            "services".to_string()
+        ];
+        let import = "...services.global.GlobalService".to_string();
+
+        let classifier = ImportClassifier::new();
+        let result = classifier.is_eligible(&import, &root_dirs);
+
+        assert_eq!(result, true);
+    }
 }

src/core/ast_parser/tree_analyzer.rs

@@ -54,8 +54,8 @@ impl TreeAnalyzer {
                         }
                     }
                     "function" => {
-                        let text = self.dot_name.get(&code, cap.node);
-                        if !text.starts_with("__") & !text.ends_with("__") {
+                        let function = self.dot_name.get(&code, cap.node);
+                        if !function.starts_with("__") & !function.ends_with("__") {
                             let mut parent = cap.node.parent();
                             let mut found_class = false;
                             while let Some(p) = parent {
@@ -64,15 +64,15 @@ impl TreeAnalyzer {
                                         let class = name_node.utf8_text(code.as_bytes())
                                         .unwrap()
                                         .to_string();
-                                        let class_method = format!("{class}.{text}");
+                                        let class_method = format!("{class}.{function}");
                                         defines.push(class_method);
                                         found_class = true;
                                     }
                                 }
                                 parent = p.parent();
                             }
                             if !found_class {
-                                defines.push(text);
+                                defines.push(function);
                             }
                         }
                     }