amql_engine/extractor/

test.rs

//! Test framework extractor (vitest, jest, mocha).
//!
//! Detects `describe()`, `it()`, `test()`, lifecycle hooks, and their
//! `.only` / `.skip` modifiers. Preserves the describe > test hierarchy
//! via nested `children`.

use super::BuiltinExtractor;
use crate::store::Annotation;
use crate::types::{AttrName, Binding, RelativePath, TagName};
use rustc_hash::FxHashMap;
use serde_json::Value as JsonValue;
use std::cell::RefCell;

/// Block names recognized as test framework calls.
const TEST_BLOCKS: &[&str] = &[
    "describe",
    "it",
    "test",
    "beforeEach",
    "afterEach",
    "beforeAll",
    "afterAll",
];

/// Modified block names (`.only`, `.skip`).
const MODIFIED_BLOCKS: &[&str] = &[
    "describe.only",
    "describe.skip",
    "it.only",
    "it.skip",
    "test.only",
    "test.skip",
];

/// Built-in test framework extractor.
pub struct TestExtractor;

impl BuiltinExtractor for TestExtractor {
    fn name(&self) -> &str {
        "test"
    }

    fn extensions(&self) -> &[&str] {
        &[".ts", ".tsx", ".js", ".jsx", ".mts", ".mjs"]
    }

    fn extract(&self, source: &str, file: &RelativePath) -> Vec<Annotation> {
        let tree = match parse_ts(source, file) {
            Some(t) => t,
            None => return vec![],
        };
        let mut annotations = Vec::new();
        visit_node(tree.root_node(), source.as_bytes(), file, &mut annotations);
        annotations
    }
}

fn visit_node(
    node: tree_sitter::Node,
    src: &[u8],
    file: &RelativePath,
    annotations: &mut Vec<Annotation>,
) {
    if node.kind() == "expression_statement" {
        if let Some(ann) = extract_test_block(node, src, file) {
            annotations.push(ann);
            return;
        }
    }

    let mut cursor = node.walk();
    for child in node.named_children(&mut cursor) {
        visit_node(child, src, file, annotations);
    }
}

/// Try to extract a test block (describe/it/test/hook) from an expression statement.
fn extract_test_block(
    node: tree_sitter::Node,
    src: &[u8],
    file: &RelativePath,
) -> Option<Annotation> {
    let call = node.named_child(0)?;
    if call.kind() != "call_expression" {
        return None;
    }

    let callee_node = call.child_by_field_name("function")?;
    let callee = node_text(callee_node, src);

    if !TEST_BLOCKS.contains(&callee.as_str()) && !MODIFIED_BLOCKS.contains(&callee.as_str()) {
        return None;
    }

    let base_name = callee.split('.').next().unwrap_or(&callee);
    let modifier = if callee.contains('.') {
        callee.split('.').nth(1)
    } else {
        None
    };

    let args = call.child_by_field_name("arguments")?;
    let args_children: Vec<_> = {
        let mut cursor = args.walk();
        args.named_children(&mut cursor).collect()
    };
    let first_arg = args_children.first();
    let label = match first_arg {
        Some(n) if n.kind() == "string" || n.kind() == "template_string" => {
            string_literal_value(*n, src)
        }
        _ => callee.clone(),
    };

    let tag = match base_name {
        "describe" => "describe",
        "it" | "test" => "test",
        _ => "hook",
    };

    let mut attrs = FxHashMap::default();
    if let Some(mod_name) = &modifier {
        attrs.insert(AttrName::from(*mod_name), JsonValue::Bool(true));
    }
    if tag == "hook" {
        attrs.insert(
            AttrName::from("kind"),
            JsonValue::String(callee.to_string()),
        );
    }

    // Extract children from callback body (second argument)
    let mut children = Vec::new();
    if args_children.len() >= 2 {
        let callback = args_children[1];
        if callback.kind() == "arrow_function" || callback.kind() == "function_expression" {
            if let Some(body) = callback.child_by_field_name("body") {
                if body.kind() == "statement_block" {
                    let mut body_cursor = body.walk();
                    for stmt in body.named_children(&mut body_cursor) {
                        visit_node(stmt, src, file, &mut children);
                    }
                }
            }
        }
    }

    Some(Annotation {
        tag: TagName::from(tag),
        attrs,
        binding: Binding::from(label),
        file: file.clone(),
        children,
    })
}

// ---------------------------------------------------------------------------
// tree-sitter helpers
// ---------------------------------------------------------------------------

fn node_text<'a>(node: tree_sitter::Node<'a>, src: &'a [u8]) -> String {
    node.utf8_text(src).unwrap_or("").to_string()
}

fn string_literal_value(node: tree_sitter::Node, src: &[u8]) -> String {
    let text = node.utf8_text(src).unwrap_or("");
    if text.len() >= 2 {
        text[1..text.len() - 1].to_string()
    } else {
        text.to_string()
    }
}

// ---------------------------------------------------------------------------
// Parser cache (thread-local)
// ---------------------------------------------------------------------------

thread_local! {
    static TS_PARSER: RefCell<Option<tree_sitter::Parser>> = const { RefCell::new(None) };
    static TSX_PARSER: RefCell<Option<tree_sitter::Parser>> = const { RefCell::new(None) };
}

fn parse_ts(source: &str, file: &RelativePath) -> Option<tree_sitter::Tree> {
    let path: &str = file.as_ref();
    let is_tsx = path.ends_with(".tsx") || path.ends_with(".jsx");

    if is_tsx {
        TSX_PARSER.with(|cell| {
            let mut opt = cell.borrow_mut();
            let parser = opt.get_or_insert_with(|| {
                let mut p = tree_sitter::Parser::new();
                p.set_language(&tree_sitter_typescript::LANGUAGE_TSX.into())
                    .expect("Failed to set TSX language");
                p
            });
            parser.parse(source, None)
        })
    } else {
        TS_PARSER.with(|cell| {
            let mut opt = cell.borrow_mut();
            let parser = opt.get_or_insert_with(|| {
                let mut p = tree_sitter::Parser::new();
                p.set_language(&tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into())
                    .expect("Failed to set TypeScript language");
                p
            });
            parser.parse(source, None)
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn run(source: &str) -> Vec<Annotation> {
        let file = RelativePath::from("src/app.test.ts");
        TestExtractor.extract(source, &file)
    }

    #[test]
    fn detects_describe_and_tests() {
        // Arrange
        let source = r#"
describe("UserService", () => {
    it("creates a user", () => {
        expect(true).toBe(true);
    });

    test("deletes a user", () => {
        expect(true).toBe(true);
    });
});
"#;

        // Act
        let anns = run(source);

        // Assert
        assert_eq!(anns.len(), 1, "should find 1 top-level describe");
        assert_eq!(anns[0].tag, "describe", "top-level should be describe");
        assert_eq!(anns[0].binding, "UserService", "describe label");
        assert_eq!(anns[0].children.len(), 2, "describe should have 2 children");
        assert_eq!(
            anns[0].children[0].tag, "test",
            "first child should be test"
        );
        assert_eq!(
            anns[0].children[0].binding, "creates a user",
            "first test label"
        );
        assert_eq!(
            anns[0].children[1].tag, "test",
            "second child should be test"
        );
        assert_eq!(
            anns[0].children[1].binding, "deletes a user",
            "second test label"
        );
    }

    #[test]
    fn detects_hooks() {
        // Arrange
        let source = r#"
describe("setup", () => {
    beforeEach(() => {
        reset();
    });

    afterAll(() => {
        cleanup();
    });

    it("works", () => {});
});
"#;

        // Act
        let anns = run(source);

        // Assert
        let desc = &anns[0];
        let hooks: Vec<_> = desc.children.iter().filter(|a| a.tag == "hook").collect();
        assert_eq!(hooks.len(), 2, "should find 2 hooks");
        assert_eq!(
            hooks[0].attrs.get("kind"),
            Some(&JsonValue::String("beforeEach".to_string())),
            "first hook kind"
        );
        assert_eq!(
            hooks[1].attrs.get("kind"),
            Some(&JsonValue::String("afterAll".to_string())),
            "second hook kind"
        );
    }

    #[test]
    fn detects_modifiers() {
        // Arrange
        let source = r#"
describe.only("focused", () => {
    it.skip("skipped test", () => {});
});
"#;

        // Act
        let anns = run(source);

        // Assert
        assert_eq!(anns.len(), 1, "should find 1 describe");
        assert_eq!(
            anns[0].attrs.get("only"),
            Some(&JsonValue::Bool(true)),
            "describe should have only modifier"
        );
        assert_eq!(
            anns[0].children[0].attrs.get("skip"),
            Some(&JsonValue::Bool(true)),
            "test should have skip modifier"
        );
    }

    #[test]
    fn nested_describes() {
        // Arrange
        let source = r#"
describe("outer", () => {
    describe("inner", () => {
        it("deep test", () => {});
    });
});
"#;

        // Act
        let anns = run(source);

        // Assert
        assert_eq!(anns.len(), 1, "should find 1 top-level describe");
        assert_eq!(
            anns[0].children.len(),
            1,
            "outer should have 1 child describe"
        );
        assert_eq!(
            anns[0].children[0].tag, "describe",
            "child should be describe"
        );
        assert_eq!(
            anns[0].children[0].children.len(),
            1,
            "inner should have 1 test"
        );
        assert_eq!(
            anns[0].children[0].children[0].binding, "deep test",
            "deep test label"
        );
    }

    #[test]
    fn detects_top_level_it() {
        // Arrange — top-level it() without describe wrapper
        let source = r#"
it("standalone test", () => {
    expect(true).toBe(true);
});
"#;

        // Act
        let anns = run(source);

        // Assert
        assert_eq!(anns.len(), 1, "should find top-level it");
        assert_eq!(anns[0].tag, "test", "should be tagged as test");
        assert_eq!(anns[0].binding, "standalone test", "test label");
    }

    #[test]
    fn detects_test_each() {
        // Arrange — test.each is not in our modified blocks, should be ignored
        let source = r#"
describe("parameterized", () => {
    test("simple test", () => {});
});
"#;

        // Act
        let anns = run(source);

        // Assert
        assert_eq!(anns.len(), 1, "should find 1 describe");
        assert_eq!(anns[0].children.len(), 1, "should have 1 test child");
    }

    #[test]
    fn detects_function_expression_callback() {
        // Arrange — using function() instead of arrow function
        let source = r#"
describe("legacy", function() {
    it("old style test", function() {
        expect(true).toBe(true);
    });
});
"#;

        // Act
        let anns = run(source);

        // Assert
        assert_eq!(anns.len(), 1, "should find 1 describe");
        assert_eq!(
            anns[0].children.len(),
            1,
            "should find test inside function expression"
        );
    }
}