orchestrata_tools.test_analyzer.body_matcher

Reference › orchestrata_tools › test_analyzer › body_matcher

BodyMatcher finds which tests call a source method by parsing the test files' bodies with the standard ast module. It was written because the serialized AST can miss calls inside fixtures, top-level functions and complex expressions. matcher uses it as its strongest strategy (body_call), and the script scripts.add_traces_from_evidence imports it. tests/test_orchestrata_tools_body_matcher.py covers it.

The division of labour is explicit. Per-node call extraction lives in call_extraction, inheritance knowledge in orchestrata_ast.inheritance, and this module owns file discovery, the fixture registry, the indexes and the query API.

• _build_index analyses each test module. _analyze_file is a pure per-file step returning a _FileAnalysis (class bases, module and class fixtures, tests with their call profiles), and _merge_analysis applies it to the shared indexes.
• FixtureRegistry stores fixture call profiles by scope and resolve looks a fixture up in pytest's order: class, module, then nearest conftest.
• The query API is direct_call_index, is_subclass_of, class_used_by_test, find_tests_for_method and is_tested.

At a glance

Class diagram

Dependency diagram

API

Classes

_FileAnalysis (dataclass)

Decorators: dataclass · Line 41

Pure per-file analysis output (phase A of _build_index).

FixtureRegistry

Line 59

Fixture profiles by scope, resolved in pytest's lookup order.

BodyMatcher

Line 98

Match source methods to tests by analyzing test body ASTs.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_tools/test_analyzer/body_matcher.py

"""Body-based test-to-source matching using Python AST.

Parses test source files directly with the `ast` module to extract
all method calls from test functions, fixtures, and helper functions.
This bypasses limitations of the serialized AST which may miss calls
in fixtures, top-level functions, or complex expressions.

Division of labour: per-node extraction lives in
:mod:`orchestrata_tools.test_analyzer.call_extraction` (pure), inheritance
knowledge in :mod:`orchestrata_ast.inheritance` (the canonical AST component); this module
owns file discovery, the fixture registry, the indexes, and the query API.
"""

from __future__ import annotations

import ast
import os
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING

from orchestrata_ast.inheritance import InheritanceMap
from orchestrata_tools.test_analyzer.call_extraction import (
    CallExtractor,
    CallProfile,
    collect_orchestrata_imports,
    get_annotation_name,
)

if TYPE_CHECKING:
    from orchestrata_tools.models.ast_models import (
        ASTCodebase,
        MethodInfo,
        ModuleInfo,
    )


@dataclass
class _FileAnalysis:
    """Pure per-file analysis output (phase A of ``_build_index``).

    Everything here derives from the file's own AST only, so files can be
    analyzed in parallel; all shared-state mutation happens in
    ``_merge_analysis``, applied strictly in discovery order.
    """

    mod_path: str
    class_bases: list[tuple[str, list[str]]] = field(default_factory=list)
    module_fixtures: dict[str, CallProfile] = field(default_factory=dict)
    class_fixtures: dict[str, dict[str, CallProfile]] = field(default_factory=dict)
    # (name, class_name, profile, decorator names, fixture params)
    tests: list[tuple[str, str | None, CallProfile, set[str], list[str]]] = field(
        default_factory=list
    )


class FixtureRegistry:
    """Fixture profiles by scope, resolved in pytest's lookup order.

    Scope keys: ``module.path`` for module fixtures,
    ``module.path::ClassName`` for class fixtures, ``pkg.conftest`` for
    conftest fixtures. Private to this module until a second consumer
    exists (see plans/bodymatcher-split-plan.md).
    """

    def __init__(self) -> None:
        self._by_scope: dict[str, dict[str, CallProfile]] = {}

    def register(self, scope: str, fixtures: dict[str, CallProfile]) -> None:
        self._by_scope[scope] = fixtures

    def resolve(
        self, mod_path: str, class_name: str | None, param: str
    ) -> CallProfile | None:
        """class fixtures -> module fixtures -> nearest-first conftests."""
        if class_name:
            found = self._by_scope.get(f"{mod_path}::{class_name}", {}).get(param)
            if found is not None:
                return found
        found = self._by_scope.get(mod_path, {}).get(param)
        if found is not None:
            return found
        # Pytest looks for fixtures in conftest.py files from the test's
        # directory up to the root tests/ directory:
        # [a, b, c] -> a.b.conftest, a.conftest, conftest
        parts = mod_path.split(".")
        for i in range(len(parts) - 1, -1, -1):
            prefix = ".".join(parts[:i]) if i > 0 else ""
            conftest_path = f"{prefix}.conftest" if prefix else "conftest"
            found = self._by_scope.get(conftest_path, {}).get(param)
            if found is not None:
                return found
        return None


class BodyMatcher:
    """Match source methods to tests by analyzing test body ASTs.

    Walks test source files directly with Python's ast module to find
    all method calls (e.g., `obj.method()`, `foo.bar.baz()`).

    Uses lightweight type inference including pytest fixture resolution:
    - Tracks variable assignments of imported classes
    - Resolves fixture return types from their bodies
    - Includes fixture calls in tests that use those fixtures
    """

    def __init__(self, test_ast: ASTCodebase, source_ast: ASTCodebase | None = None):
        self.test_ast = test_ast
        self.source_ast = source_ast
        self._extractor = CallExtractor()
        # method_name -> list of (test_module, test_name, test_class)
        self._calls: dict[str, list[tuple[str, str, str | None]]] = defaultdict(list)
        # last segment of a dotted call name -> list of call_names sharing it,
        # e.g. "run" -> ["Foo.run", "Bar.Baz.run"]. Built once in _build_index
        # after self._calls is fully populated; lets find_tests_for_method's
        # suffix-matching patterns look up a shortlist instead of scanning
        # every entry in self._calls.
        self._calls_by_suffix: dict[str, list[str]] = defaultdict(list)
        # (test_module, test_name) -> set of imported class names used in that test
        self._test_imports: dict[tuple[str, str], set[str]] = {}
        # module -> union of imported class names across the module's tests;
        # class_used_by_test needs only per-module existence, so this union
        # replaces a full scan of self._test_imports per call.
        self._imports_by_module: dict[str, set[str]] = {}
        self._fixtures = FixtureRegistry()
        # class -> ancestors, from the source AST + test-file class defs
        self._inheritance = InheritanceMap.from_codebase(source_ast)
        self._parsed_modules: set[str] = set()
        # (test_module, test_name, test_class) -> (direct body calls,
        # imported classes used directly) — captured BEFORE fixture merging,
        # so fixture-side calls never leak into per-test attribution.
        self._direct_calls: dict[
            tuple[str, str, str | None], tuple[set[str], set[str]]
        ] = {}
        self._build_index()

    def _build_index(self) -> None:
        """Pre-parse all test modules to build the call index.

        Discovers test files both from the serialized AST and by scanning
        the tests/ directory directly. This catches newer test files that
        may not have been re-serialized yet.
        """
        # Discover files in the exact legacy order — AST modules first, then
        # the tests/ directory scan, then conftest.py files — with the same
        # dedupe rules (by mod_path; by file path for the directory scan).
        files: list[tuple[str, str]] = []

        def enqueue(source_path: str, mod_path: str) -> None:
            if mod_path in self._parsed_modules:
                return
            self._parsed_modules.add(mod_path)
            files.append((source_path, mod_path))

        ast_paths: set[str] = set()
        for module in self.test_ast.modules:
            source_path = module.file_path
            if source_path and Path(source_path).exists():
                ast_paths.add(source_path)
                enqueue(source_path, module.path)

        # Also scan tests/ directory for any files not in the AST.
        # Derive root from file paths since root_path may be "tests" itself.
        test_root: Path | None = None
        for module in self.test_ast.modules:
            fp = module.file_path
            if fp and ("tests/" in fp or fp.startswith("tests")):
                path_parts = Path(fp).parts
                for i, part in enumerate(path_parts):
                    if part == "tests":
                        test_root = Path(*path_parts[: i + 1])
                        break
                if test_root is not None:
                    break
        if test_root is None:
            test_root = Path(self.test_ast.root_path)
            if test_root.name == "tests":
                test_root = test_root.parent / "tests"
            else:
                test_root = test_root / "tests"

        if test_root.exists():
            for py_file in test_root.rglob("*.py"):
                if py_file.name.startswith("_"):
                    continue
                file_path = str(py_file.resolve())
                if file_path not in ast_paths:
                    rel_path = py_file.relative_to(test_root)
                    mod_path = str(rel_path.with_suffix("")).replace("/", ".")
                    enqueue(file_path, mod_path)

        for conftest in test_root.rglob("conftest.py"):
            file_path = str(conftest.resolve())
            rel_path = conftest.relative_to(test_root)
            mod_path = str(rel_path.with_suffix("")).replace("/", ".")
            enqueue(file_path, mod_path)

        # Phase A (parallel): pure per-file analysis — Python 3.14t is
        # free-threaded, so CPU-bound parsing scales across threads.
        # Phase B (sequential): executor.map yields in submission order, so
        # merging in this loop replays the exact legacy discovery order.
        if files:
            with ThreadPoolExecutor(
                max_workers=min(8, os.cpu_count() or 1)
            ) as executor:
                for analysis in executor.map(
                    lambda item: self._analyze_file(*item), files
                ):
                    if analysis is not None:
                        self._merge_analysis(analysis)

        # self._calls is fully populated now — build the suffix shortlist
        # index once so find_tests_for_method's suffix-matching patterns
        # don't have to scan every entry in self._calls per lookup.
        for call_name in self._calls:
            parts = call_name.split(".")
            if len(parts) >= 2:
                self._calls_by_suffix[parts[-1]].append(call_name)

        # self._test_imports is fully populated now — build the per-module
        # import union used by class_used_by_test.
        for (mod, _test), imports in self._test_imports.items():
            self._imports_by_module.setdefault(mod, set()).update(imports)

    def _parse_module(self, module: ModuleInfo) -> None:
        """Parse a single test module and extract calls."""
        mod_path = module.path
        source_path = module.file_path
        if not source_path:
            return
        self._parse_file(source_path, mod_path)

    def _parse_file(self, source_path: str, mod_path: str) -> None:
        """Parse a single Python test file and extract calls (serial path)."""
        if mod_path in self._parsed_modules:
            return
        self._parsed_modules.add(mod_path)
        analysis = self._analyze_file(source_path, mod_path)
        if analysis is not None:
            self._merge_analysis(analysis)

    def _analyze_file(self, source_path: str, mod_path: str) -> _FileAnalysis | None:
        """Parse one file and extract its analysis.

        Pure — no shared-state writes — so it is safe to run in parallel
        across files (the CallExtractor is stateless).
        """
        try:
            with open(source_path, encoding="utf-8") as f:
                tree = ast.parse(f.read(), filename=source_path)
        except SyntaxError, UnicodeDecodeError, FileNotFoundError:
            # FileNotFoundError: rglob enumeration races concurrent scratch-dir
            # teardown (parallel test workers); a vanished file is equivalent
            # to one that was never enumerated.
            return None

        analysis = _FileAnalysis(mod_path=mod_path)
        imported_names, imported_classes = collect_orchestrata_imports(tree)

        # Collect class definitions and inheritance from test files
        local_classes: set[str] = set()
        for node in tree.body:
            if isinstance(node, ast.ClassDef):
                local_classes.add(node.name)
                bases = [
                    name for base in node.bases if (name := get_annotation_name(base))
                ]
                analysis.class_bases.append((node.name, bases))

        def profile_of(fn: ast.FunctionDef | ast.AsyncFunctionDef) -> CallProfile:
            return self._extractor.extract(
                fn, imported_names, imported_classes, local_classes
            )

        # Phase 1: Collect fixtures and their profiles
        for node in tree.body:
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                if self._extractor.is_pytest_fixture(node):
                    analysis.module_fixtures[node.name] = profile_of(node)
            elif isinstance(node, ast.ClassDef):
                for child in node.body:
                    if isinstance(
                        child, (ast.FunctionDef, ast.AsyncFunctionDef)
                    ) and self._extractor.is_pytest_fixture(child):
                        analysis.class_fixtures.setdefault(node.name, {})[
                            child.name
                        ] = profile_of(child)

        # Phase 2: Collect test profiles + pure per-test metadata
        def add_test(
            fn: ast.FunctionDef | ast.AsyncFunctionDef, class_name: str | None
        ) -> None:
            analysis.tests.append(
                (
                    fn.name,
                    class_name,
                    profile_of(fn),
                    self._extractor.decorator_names(fn),
                    self._extractor.fixture_params(fn),
                )
            )

        for node in tree.body:
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                if node.name.startswith("test_"):
                    add_test(node, None)
            elif isinstance(node, ast.ClassDef):
                for child in node.body:
                    if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
                        if child.name.startswith("test_"):
                            add_test(child, node.name)
        return analysis

    def _merge_analysis(self, analysis: _FileAnalysis) -> None:
        """Apply one file's analysis to the shared indexes.

        Must run sequentially in discovery order: fixture registration for a
        file happens only after all earlier files were fully indexed (tests
        parsed before their conftest deliberately do not see its fixtures),
        and ``_calls`` list order is observable matching behavior.
        """
        mod_path = analysis.mod_path
        for cls_name, bases in analysis.class_bases:
            self._inheritance.add(cls_name, bases)

        self._fixtures.register(mod_path, analysis.module_fixtures)
        for cls_name, fixtures in analysis.class_fixtures.items():
            self._fixtures.register(f"{mod_path}::{cls_name}", fixtures)

        for name, class_name, profile, dec_names, fixture_params in analysis.tests:
            # A test's own decorators are metadata, not exercised behaviour.
            self._direct_calls[(mod_path, name, class_name)] = (
                set(profile.calls) - dec_names,
                set(profile.used_imports),
            )
            # Merge fixture profiles in pytest's lookup order
            for param in fixture_params:
                fixture = self._fixtures.resolve(mod_path, class_name, param)
                if fixture is not None:
                    profile = profile.merged_with(fixture)
            self._test_imports[(mod_path, name)] = set(profile.used_imports)
            for call_name in profile.calls:
                self._calls[call_name].append((mod_path, name, class_name))

    def direct_call_index(
        self,
    ) -> dict[tuple[str, str, str | None], tuple[set[str], set[str]]]:
        """Per test: (direct body calls, imported classes used directly).

        Keys are ``(test_module, test_name, test_class)``. Calls are the
        var-type-resolved names from the test's OWN body (``obj.method()``
        with ``obj = ImportedClass()`` resolves to ``ImportedClass.method``);
        fixture-resolved calls are deliberately excluded — a test's fixtures
        exercise infrastructure, not the behaviour the test attributes.
        Callers receive fresh copies.
        """
        return {k: (set(c), set(i)) for k, (c, i) in self._direct_calls.items()}

    def is_subclass_of(self, child: str, parent: str) -> bool:
        """Public inheritance check over the source AST (child == parent counts)."""
        return self._inheritance.is_subclass(child, parent)

    def class_used_by_test(self, test_module: str, class_name: str) -> bool:
        """True if any test in *test_module* uses *class_name* or a subclass.

        Backed by direct source parsing, so it catches imports the
        serialized AST may have missed (TestMatcher's fallback path).
        Served from the per-module import union built in ``_build_index``.
        """
        imports = self._imports_by_module.get(test_module)
        if not imports:
            return False
        if class_name in imports:
            return True
        for imp_class in imports:
            if self._inheritance.is_subclass(imp_class, class_name):
                return True
        return False

    def find_tests_for_method(
        self, method: MethodInfo
    ) -> list[tuple[str, str, str | None]]:
        """Find all tests that call a given source method.

        Matches patterns like:
        - ClassName.method_name (direct class method reference)
        - obj.method_name (resolved via type inference)
        - method_name (bare function call)
        """
        results: list[tuple[str, str, str | None]] = []
        seen: set[tuple[str, str, str | None]] = set()

        def _add(tests: list[tuple[str, str, str | None]]) -> None:
            for t in tests:
                if t not in seen:
                    seen.add(t)
                    results.append(t)

        # Pattern 1: "ClassName.method_name" (exact match)
        class_method = f"{method.class_name}.{method.name}"
        _add(self._calls.get(class_method, []))

        # Pattern 1b: subclass calls — "Subclass.method_name" tests parent method
        for call_name in self._calls_by_suffix.get(method.name, ()):
            parts = call_name.split(".")
            caller_class = parts[0] if len(parts) == 2 else ".".join(parts[:-1])
            if self._inheritance.is_subclass(caller_class, method.class_name):
                _add(self._calls[call_name])

        # Pattern 2: "obj.method_name" — any call ending with .method_name
        # where the test imports the source class
        for call_name in self._calls_by_suffix.get(method.name, ()):
            for mod_path, test_name, test_cls in self._calls[call_name]:
                test_key = (mod_path, test_name)
                if test_key in self._test_imports:
                    if method.class_name in self._test_imports[test_key]:
                        _add([(mod_path, test_name, test_cls)])

        # Pattern 3: bare "method_name"
        _add(self._calls.get(method.name, []))

        return results

    def is_tested(self, method: MethodInfo, test_ast: ASTCodebase) -> bool:
        """Check if a source method is called in any test."""
        tests = self.find_tests_for_method(method)
        if not tests:
            return False
        # For unqualified calls, verify the test imports the source class
        for mod_path, test_name, _test_class in tests:
            test_key = (mod_path, test_name)
            if test_key in self._test_imports:
                imports = self._test_imports[test_key]
                if method.class_name in imports:
                    return True
                # Also accept if a subclass is imported
                for imp_class in imports:
                    if self._inheritance.is_subclass(imp_class, method.class_name):
                        return True
            # Fallback: check module-level imports
            for mod in test_ast.modules:
                if mod.path == mod_path:
                    for imp in mod.imports:
                        if method.class_name in imp.names:
                            return True
                    break
        return False