amql_engine/extractor/

go_test.rs

//! Go test and benchmark extractor.
//!
//! Detects `func Test*(t *testing.T)` and `func Benchmark*(b *testing.B)`
//! patterns in Go source files. Uses tree-sitter-go for parsing.

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;

/// Built-in Go test/benchmark extractor.
pub struct GoTestExtractor;

impl BuiltinExtractor for GoTestExtractor {
    fn name(&self) -> &str {
        "go_test"
    }

    fn extensions(&self) -> &[&str] {
        &[".go"]
    }

    fn extract(&self, source: &str, file: &RelativePath) -> Vec<Annotation> {
        // Only process _test.go files
        let path: &str = file.as_ref();
        if !path.ends_with("_test.go") {
            return vec![];
        }

        let tree = match parse_go(source) {
            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() == "function_declaration" {
        if let Some(ann) = extract_test_func(node, src, file) {
            annotations.push(ann);
        }
    }

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

/// Detect `func TestFoo(t *testing.T)` or `func BenchmarkFoo(b *testing.B)`.
fn extract_test_func(
    node: tree_sitter::Node,
    src: &[u8],
    file: &RelativePath,
) -> Option<Annotation> {
    let name_node = node.child_by_field_name("name")?;
    let name = node_text(&name_node, src);

    let (tag, prefix) = if name.starts_with("Test") {
        ("test", "Test")
    } else if name.starts_with("Benchmark") {
        ("benchmark", "Benchmark")
    } else {
        return None;
    };

    // Validate the parameter type: *testing.T or *testing.B
    let params = node.child_by_field_name("parameters")?;
    let param_text = node_text(&params, src);
    let expected_type = if tag == "test" {
        "testing.T"
    } else {
        "testing.B"
    };
    if !param_text.contains(expected_type) {
        return None;
    }

    let test_name = name.strip_prefix(prefix).unwrap_or(&name);

    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        JsonValue::String(test_name.to_string()),
    );

    Some(Annotation {
        tag: TagName::from(tag),
        attrs,
        binding: Binding::from(name),
        file: file.clone(),
        children: vec![],
    })
}

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

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

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

fn parse_go(source: &str) -> Option<tree_sitter::Tree> {
    GO_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_go::LANGUAGE.into())
                .expect("Failed to set Go language for tree-sitter");
            p
        });
        parser.parse(source, None)
    })
}

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

    fn run(source: &str) -> Vec<Annotation> {
        let file = RelativePath::from("pkg/handler_test.go");
        GoTestExtractor.extract(source, &file)
    }

    #[test]
    fn detects_test_functions() {
        // Arrange
        let source = r#"package handler

import "testing"

func TestCreateUser(t *testing.T) {
	// test body
}

func TestDeleteUser(t *testing.T) {
	// test body
}
"#;

        // Act
        let anns = run(source);

        // Assert
        let tests: Vec<_> = anns.iter().filter(|a| a.tag == "test").collect();
        assert_eq!(tests.len(), 2, "should find 2 test functions");
        assert_eq!(tests[0].binding, "TestCreateUser", "first test binding");
        assert_eq!(tests[1].binding, "TestDeleteUser", "second test binding");
        assert_eq!(
            tests[0].attrs.get("name"),
            Some(&JsonValue::String("CreateUser".to_string())),
            "test name without prefix"
        );
    }

    #[test]
    fn detects_benchmark_functions() {
        // Arrange
        let source = r#"package handler

import "testing"

func BenchmarkSerialize(b *testing.B) {
	for i := 0; i < b.N; i++ {
		serialize()
	}
}
"#;

        // Act
        let anns = run(source);

        // Assert
        let benchmarks: Vec<_> = anns.iter().filter(|a| a.tag == "benchmark").collect();
        assert_eq!(benchmarks.len(), 1, "should find 1 benchmark");
        assert_eq!(
            benchmarks[0].binding, "BenchmarkSerialize",
            "benchmark binding"
        );
        assert_eq!(
            benchmarks[0].attrs.get("name"),
            Some(&JsonValue::String("Serialize".to_string())),
            "benchmark name without prefix"
        );
    }

    #[test]
    fn ignores_non_test_files() {
        // Arrange
        let source = r#"package handler

import "testing"

func TestCreateUser(t *testing.T) {}
"#;
        let file = RelativePath::from("pkg/handler.go");

        // Act
        let anns = GoTestExtractor.extract(source, &file);

        // Assert
        assert!(anns.is_empty(), "should ignore non-_test.go files");
    }

    #[test]
    fn ignores_helper_functions() {
        // Arrange
        let source = r#"package handler

import "testing"

func TestMain(m *testing.M) {
	os.Exit(m.Run())
}

func helperSetup(t *testing.T) {
	// not a test
}
"#;

        // Act
        let anns = run(source);

        // Assert
        assert!(
            anns.is_empty(),
            "should ignore TestMain (testing.M) and helpers"
        );
    }

    #[test]
    fn detects_mixed_tests_and_benchmarks() {
        // Arrange
        let source = r#"package handler

import "testing"

func TestParse(t *testing.T) {}
func BenchmarkParse(b *testing.B) {}
func TestFormat(t *testing.T) {}
"#;

        // Act
        let anns = run(source);

        // Assert
        let tests: Vec<_> = anns.iter().filter(|a| a.tag == "test").collect();
        let benchmarks: Vec<_> = anns.iter().filter(|a| a.tag == "benchmark").collect();
        assert_eq!(tests.len(), 2, "should find 2 tests");
        assert_eq!(benchmarks.len(), 1, "should find 1 benchmark");
    }
}