orchestrata_tools.test_generator.testgen.generator

Reference › orchestrata_tools › test_generator › testgen › generator

This module is the top-level orchestrator of the testgen pipeline. It is re-exported by the testgen package, and tests/test_orchestrata_tools_testgen_generator_coverage.py and ..._units.py cover it. It works from a CodebaseAST and the already-parsed trees of the same analyzer run, from orchestrata_ast, and never re-parses source.

generate_tests runs the pipeline: rank functions by priority with ranking, generate structural tests for dataclasses, enums, abstract classes and protocols with structural, generate branch tests with branch, and write the files. Its parameters select the output directory (tests/generated by default), target modules, how many ranked entry points to stub, public-only symbols, and one file per module or a single combined file, plus a ProjectConfig (config). Ranking runs over the whole snapshot so fan-in reflects all modules. A module is skipped if it has no classes or functions, has no parsed tree, or its import path is not importable. The file name comes from the config pattern. It returns a dictionary of output path to rendered source. print_ranking prints the priority table before generation. The classes are located in trees through adapter.

At a glance

Class diagram

No classes; see Functions.

Dependency diagram

API

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_tools/test_generator/testgen/generator.py

"""generator.py — Top-level orchestrator.

Drives the full Layer 1 + 2 pipeline:
  1. Rank entry points by fan_in / reachable
  2. Generate structural tests (dataclass / enum / abstract / protocol)
  3. Generate branch stubs for top-ranked functions
  4. Write one pytest file per module (or a single combined file)

Usage
-----
    from orchestrata_ast.analyzer import CodebaseAnalyzer
    from orchestrata_tools.test_generator.testgen import generate_tests
    from orchestrata_tools.test_generator.testgen.config import ORCHESTRATA_ULTIMATE

    analyzer = CodebaseAnalyzer("orchestrata-ultimate", Path("src/orchestrata_ultimate"))
    codebase = analyzer.analyze()
    results = generate_tests(codebase, analyzer.trees, out_dir="tests/generated",
                             project_config=ORCHESTRATA_ULTIMATE)

Restored 2026-08-19 from ``scripts/ast_serializer/testgen`` (deleted
2026-07-09, commit 6a186c9d) and re-pointed at the current AST layer: the
deleted ``ProjectRecord``/``CallGraphAnalyzer`` inputs are replaced by
``orchestrata_ast.models.CodebaseAST`` plus the already-parsed ``ast``
trees ``orchestrata_ast.analyzer.CodebaseAnalyzer`` produces while
building it (``.trees: dict[file_path, ast.Module]``) — see
``orchestrata_tools.test_generator.testgen.adapter`` for how individual
class/function nodes are located in those trees.
"""

from __future__ import annotations

import ast
import importlib
from pathlib import Path
from typing import TYPE_CHECKING

from orchestrata_tools.test_generator.testgen.adapter import (
    class_structural_extra,
    find_class_node,
    find_function_node,
    looks_abstract_via_mixin,
)
from orchestrata_tools.test_generator.testgen.branch import gen_branch_tests
from orchestrata_tools.test_generator.testgen.config import ProjectConfig
from orchestrata_tools.test_generator.testgen.ranking import (
    RankedEntry,
    rank_functions,
    summarise,
)
from orchestrata_tools.test_generator.testgen.structural import (
    gen_abstract_tests,
    gen_dataclass_tests,
    gen_entry_point_tests,
    gen_enum_tests,
    gen_protocol_tests,
)
from orchestrata_tools.test_generator.testgen.writer import TestFile

if TYPE_CHECKING:
    from orchestrata_ast.models import ClassInfo, CodebaseAST, ModuleInfo

# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def generate_tests(
    codebase: CodebaseAST,
    trees: dict[str, ast.Module],
    out_dir: str | Path = "tests/generated",
    target_modules: list[str] | None = None,
    top_entry_points: int = 10,
    public_only: bool = True,
    per_module: bool = True,
    project_config: ProjectConfig | None = None,
) -> dict[str, str]:
    """
    Run the full Layer 1 + 2 pipeline and write test files.

    Parameters
    ----------
    codebase            : a serialized CodebaseAST (orchestrata_ast.models),
                          from CodebaseAnalyzer(...).analyze()
    trees               : file_path -> ast.Module, from the SAME analyzer
                          run (``CodebaseAnalyzer.trees``) — this generator
                          never re-parses source.
    out_dir             : directory to write test files into
    target_modules      : dotted module paths (``ModuleInfo.path``) to
                          generate for; None = every module in *codebase*
    top_entry_points    : how many ranked entry points to stub per module
    public_only         : skip private symbols
    per_module          : one file per source module (False = single file)
    project_config      : project-specific configuration (package prefix,
                          import allow-list, etc.)

    Returns
    -------
    dict mapping output file path → rendered source
    """
    config = project_config or ProjectConfig()
    out_path = Path(out_dir)
    out_path.mkdir(parents=True, exist_ok=True)

    modules = [
        m
        for m in codebase.modules
        if target_modules is None or m.path in target_modules
    ]

    # Rank across every module in scope, so fan_in/reachable reflect the
    # whole snapshot, not just the modules being written out this run.
    all_ranked = rank_functions(codebase, public_only=public_only)

    results: dict[str, str] = {}

    if per_module:
        for mod in modules:
            if not mod.classes and not mod.functions:
                continue
            tree = trees.get(mod.file_path)
            if tree is None:
                continue
            import_path = _resolve_import_path(mod.path, config)
            if not _module_importable(import_path):
                continue
            ranked_for_mod = [e for e in all_ranked if e.module == mod.path]
            src = _generate_module_file(
                mod, tree, ranked_for_mod, codebase, top_entry_points, config
            )
            if src.strip():
                fname_pat = config.test_file_pattern.replace(
                    "{mod_name}", mod.path.replace(".", "_")
                )
                fname = out_path / fname_pat
                fname.write_text(src, encoding="utf-8")
                results[str(fname)] = src
    else:
        src = _generate_combined_file(
            modules, trees, all_ranked, codebase, top_entry_points, config
        )
        fname = out_path / "test_generated.py"
        fname.write_text(src, encoding="utf-8")
        results[str(fname)] = src

    return results


def print_ranking(codebase: CodebaseAST, top: int = 30) -> None:
    """Print a priority table to stdout — useful before generating."""
    ranked = rank_functions(codebase)
    print(summarise(ranked, top=top))


# ---------------------------------------------------------------------------
# Per-module file generation
# ---------------------------------------------------------------------------


def _module_importable(module_path: str) -> bool:
    """Check if a module can be imported without raising."""
    try:
        importlib.import_module(module_path)
        return True
    except Exception:
        return False


def _resolve_import_path(mod_path: str, config: ProjectConfig) -> str:
    """Return the full dotted import path for a module."""
    if config.package_prefix:
        return (
            f"{config.package_prefix}.{mod_path}" if mod_path else config.package_prefix
        )
    return mod_path


def _generate_module_file(
    mod: ModuleInfo,
    tree: ast.Module,
    ranked: list[RankedEntry],
    codebase: CodebaseAST,
    top_n: int,
    config: ProjectConfig,
) -> str:
    tf = TestFile(module_under_test=mod.path, config=config)

    if config.testfile_hook is not None:
        config.testfile_hook(tf)

    import_path = _resolve_import_path(mod.path, config)
    names = _public_names(mod, import_path)
    tf.import_lines.append(f"from {import_path} import (")
    for nm in names:
        tf.import_lines.append(f"    {nm},")
    tf.import_lines.append(")")

    sections_added = 0

    # ---- Structural tests (dataclass / enum / abstract / protocol) ----
    for cls in mod.classes:
        if not cls.name or cls.name.startswith("_"):
            continue
        if not (cls.is_dataclass or cls.is_enum or cls.is_protocol or cls.is_abstract):
            continue
        class_node = find_class_node(cls.name, cls.line, tree)
        if class_node is None:
            continue
        extra = class_structural_extra(cls, class_node)

        if extra["is_dataclass"]:
            tf.add_section(
                gen_dataclass_tests(cls.name, extra, import_path, tf, config)
            )
            sections_added += 1

        if extra["is_enum"]:
            tf.add_section(gen_enum_tests(cls.name, extra, import_path, tf, config))
            sections_added += 1

        if (
            extra["is_abstract"]
            and not extra["is_protocol"]
            and extra["abstract_methods"]
        ):
            tf.add_section(gen_abstract_tests(cls.name, extra, import_path, tf, config))
            sections_added += 1

        if extra["is_protocol"]:
            tf.add_section(gen_protocol_tests(cls.name, extra, import_path, tf, config))
            sections_added += 1

    # ---- Entry point stubs ----
    if config.enable_entry_point_stubs and ranked:
        ep_block = gen_entry_point_tests(ranked, top_n=top_n)
        if ep_block.strip():
            tf.add_section(ep_block)
            sections_added += 1

    # ---- Branch stubs for top ranked functions ----
    if config.enable_branch_tests:
        try:
            file_text = Path(mod.file_path).read_text(encoding="utf-8")
        except OSError:
            file_text = ""
        class_by_name: dict[str, ClassInfo] = {c.name: c for c in mod.classes}

        for entry in ranked[:top_n]:
            func_node = find_function_node(entry.symbol.name, entry.symbol.line, tree)
            if func_node is None:
                continue

            class_node = None
            looks_abstract = False
            if entry.is_method and entry.class_name:
                cls_info = class_by_name.get(entry.class_name)
                if cls_info is not None:
                    class_node = find_class_node(cls_info.name, cls_info.line, tree)
                    looks_abstract = looks_abstract_via_mixin(cls_info, codebase)

            source = (
                ast.get_source_segment(file_text, func_node) or "" if file_text else ""
            )
            branch_src = gen_branch_tests(
                func_name=entry.symbol.name,
                func_qname=entry.qualified_name,
                func_node=func_node,
                source=source,
                is_method=entry.is_method,
                class_name=entry.class_name,
                class_node=class_node,
                tf=tf,
                config=config,
                looks_abstract_via_mixin=looks_abstract,
            )
            if branch_src.strip():
                tf.add_section(branch_src)
                sections_added += 1

    if sections_added == 0:
        return ""

    return tf.render()


def _generate_combined_file(
    modules: list[ModuleInfo],
    trees: dict[str, ast.Module],
    all_ranked: list[RankedEntry],
    codebase: CodebaseAST,
    top_n: int,
    config: ProjectConfig,
) -> str:
    tf = TestFile(module_under_test="(all modules)", config=config)

    for mod in modules:
        import_path = _resolve_import_path(mod.path, config)
        names = _public_names(mod, import_path)
        if names:
            tf.import_lines.append(f"# --- {mod.path} ---")
            tf.import_lines.append(f"from {import_path} import (")
            for nm in names:
                tf.import_lines.append(f"    {nm},")
            tf.import_lines.append(")")

    for mod in modules:
        tree = trees.get(mod.file_path)
        if tree is None:
            continue
        ranked_for_mod = [e for e in all_ranked if e.module == mod.path]
        src = _generate_module_file(mod, tree, ranked_for_mod, codebase, top_n, config)
        if src.strip():
            # Strip the file header — we're combining
            body = "\n".join(src.splitlines()[8:])  # skip docstring+imports
            tf.add_section(f"# {'=' * 68}\n# Module: {mod.path}\n# {'=' * 68}\n{body}")

    return tf.render()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _public_names(mod: ModuleInfo, import_path: str) -> list[str]:
    """Names to import into the generated test file's header.

    Bug found while restoring: this used to return JUST ``mod.all_exports``
    when a module defines ``__all__``, treating it as a restrictive filter.
    But the structural-test loop in ``_generate_module_file`` iterates
    every non-private class regardless of ``__all__`` — so any
    dataclass/enum/protocol/abstract class NOT listed in ``__all__`` still
    got a full ``TestXBranches``/``TestXDataclass`` section generated
    against it, referencing the bare class name, while the header import
    silently dropped that same name — a generated file that fails with
    ``NameError`` on every test for that class the moment it's run.
    Confirmed live: a two-class module with ``__all__ = ['Shown']`` produced
    3 failing/3 passing tests instead of 6 passing. Fixed by treating
    ``__all__`` as ADDITIVE to the AST-derived names (widen the import
    list, never narrow it below what generation actually references) —
    this only ever adds names a hand-written re-export might need, it
    never drops a name the generated tests depend on.
    """
    names = [c.name for c in mod.classes if not c.name.startswith("_")]
    names += [f.name for f in mod.functions if not f.name.startswith("_")]
    if mod.all_exports:
        names = sorted(set(names) | set(mod.all_exports))

    # Verify names are actually importable at runtime (best-effort)
    try:
        imported = importlib.import_module(import_path)
    except Exception:
        # Module fails to import (e.g. click callback errors),
        # fall back to AST-derived names
        return names
    return [n for n in names if hasattr(imported, n)]