archlens/crates/adapters/tree-sitter/tests/python_analyzer_tests.rs

use archlens_domain::{
    CodeElementKind, FilePath, Language, RelationshipKind, SourceFile, ports::SourceAnalyzer,
};
use archlens_tree_sitter::TreeSitterAnalyzer;

fn analyze_python(source: &str, filename: &str) -> archlens_domain::AnalysisResult {
    let dir = tempfile::tempdir().unwrap();
    let file_path = dir.path().join(filename);
    std::fs::write(&file_path, source).unwrap();

    let analyzer = TreeSitterAnalyzer::new();
    let source_file = SourceFile::new(
        FilePath::new(file_path.to_str().unwrap()).unwrap(),
        Language::Python,
    );
    analyzer.analyze_file(&source_file).unwrap()
}

#[test]
fn extracts_python_class() {
    let result = analyze_python("class Order:\n    pass\n", "order.py");

    assert_eq!(result.elements().len(), 1);
    assert_eq!(result.elements()[0].name(), "Order");
    assert_eq!(result.elements()[0].kind(), CodeElementKind::Class);
}

#[test]
fn extracts_python_inheritance() {
    let source = "class Animal:\n    pass\n\nclass Dog(Animal):\n    pass\n";
    let result = analyze_python(source, "animals.py");

    assert_eq!(result.elements().len(), 2);

    let inheritance: Vec<_> = result
        .relationships()
        .iter()
        .filter(|r| r.kind() == RelationshipKind::Inheritance)
        .collect();

    assert_eq!(inheritance.len(), 1);
    assert_eq!(inheritance[0].source(), "Dog");
    assert_eq!(inheritance[0].target(), "Animal");
}

#[test]
fn extracts_composition_from_type_annotated_fields() {
    let source = "class Address:\n    pass\n\nclass User:\n    def __init__(self):\n        self.address: Address = Address()\n";
    let result = analyze_python(source, "user.py");

    let composition: Vec<_> = result
        .relationships()
        .iter()
        .filter(|r| r.kind() == RelationshipKind::Composition)
        .collect();

    assert_eq!(composition.len(), 1);
    assert_eq!(composition[0].source(), "User");
    assert_eq!(composition[0].target(), "Address");
}

#[test]
fn extracts_import_from_import_statement() {
    let source = "import os\nfrom commons.src.schema import BaseModel\n";
    let result = analyze_python(source, "service.py");

    let imports: Vec<_> = result
        .relationships()
        .iter()
        .filter(|r| r.kind() == RelationshipKind::Import)
        .collect();

    assert!(imports.iter().any(|r| r.target() == "commons.src.schema"));
    assert!(
        !imports.iter().any(|r| r.target() == "os"),
        "stdlib should be filtered"
    );
}

#[test]
fn extracts_relative_imports_from_init() {
    let source = "from .schema import BaseModel\nfrom .client import ApiClient\n";
    let result = analyze_python(source, "__init__.py");

    let imports: Vec<_> = result
        .relationships()
        .iter()
        .filter(|r| r.kind() == RelationshipKind::Import)
        .collect();

    assert_eq!(imports.len(), 2);
}

#[test]
fn extracts_import_from_plain_import() {
    let source = "import commons.utils\n";
    let result = analyze_python(source, "service.py");

    let imports: Vec<_> = result
        .relationships()
        .iter()
        .filter(|r| r.kind() == RelationshipKind::Import)
        .collect();

    assert!(imports.iter().any(|r| r.target() == "commons.utils"));
}

#[test]
fn extracts_composition_from_constructor_params() {
    let source = "class Config:\n    pass\n\nclass Service:\n    def __init__(self, config: Config):\n        pass\n";
    let result = analyze_python(source, "service.py");

    let composition: Vec<_> = result
        .relationships()
        .iter()
        .filter(|r| r.kind() == RelationshipKind::Composition)
        .collect();

    assert_eq!(composition.len(), 1);
    assert_eq!(composition[0].source(), "Service");
    assert_eq!(composition[0].target(), "Config");
}

#[test]
fn extracts_composition_from_class_level_annotations() {
    let source = "class Gad:\n    pass\n\nclass Definition:\n    gad: Gad\n    name: str\n";
    let result = analyze_python(source, "models.py");

    let composition: Vec<_> = result
        .relationships()
        .iter()
        .filter(|r| r.kind() == RelationshipKind::Composition)
        .collect();

    assert_eq!(composition.len(), 1);
    assert_eq!(composition[0].source(), "Definition");
    assert_eq!(composition[0].target(), "Gad");
}

#[test]
fn extracts_python_class_methods() {
    let source = "class OrderService:\n    def process(self):\n        pass\n    def cancel(self):\n        pass\n";
    let result = analyze_python(source, "service.py");

    let element = result
        .elements()
        .iter()
        .find(|e| e.name() == "OrderService")
        .unwrap();

    assert!(
        element.methods().iter().any(|m| m.contains("process")),
        "expected 'process' method, got: {:?}",
        element.methods()
    );
    assert!(
        element.methods().iter().any(|m| m.contains("cancel")),
        "expected 'cancel' method, got: {:?}",
        element.methods()
    );
}

#[test]
fn extracts_python_method_typed_params() {
    let source = "class OrderService:\n    def process(self, order: Order, count: int) -> None:\n        pass\n";
    let result = analyze_python(source, "service.py");

    let element = result
        .elements()
        .iter()
        .find(|e| e.name() == "OrderService")
        .unwrap();

    let method = element
        .methods()
        .iter()
        .find(|m| m.contains("process"))
        .unwrap();
    assert!(
        method.contains("order: Order"),
        "missing typed param: {method}"
    );
    assert!(
        method.contains("count: int"),
        "missing typed param: {method}"
    );
}

#[test]
fn extracts_python_method_return_annotation() {
    let source = "class OrderService:\n    def get(self) -> Order:\n        pass\n";
    let result = analyze_python(source, "service.py");

    let element = result
        .elements()
        .iter()
        .find(|e| e.name() == "OrderService")
        .unwrap();

    let method = element
        .methods()
        .iter()
        .find(|m| m.contains("get"))
        .unwrap();
    assert!(
        method.contains("-> Order"),
        "expected return type, got: {method}"
    );
}