orchestrata_tools.test_generator.testgen.branch

Reference › orchestrata_tools › test_generator › testgen › branch

This module generates real, runnable branch and exception tests from a function's AST. For every branch event it emits a complete pytest test and not a stub, deriving inputs from the condition. An if produces two tests (true path and false path), a raise produces one test that reaches it under pytest.raises, a try produces one test per handler, and an assert produces a passing and a violating test. Inputs come from conditions. Where an input type cannot be resolved the test uses cast(Any, None) with the typing import registered on the output file, never a mock. The generated tests contain no skip, pass or TODO. It was restored in August 2026 and walks the already-parsed function node directly (see adapter).

gen_branch_tests is the single public function. It takes the function name, qualified name, function node, source text, whether it is a method, the class name and node, an optional TestFile (writer), a ProjectConfig (config) and a flag for classes made abstract by a mixin. It collects branch events, and when there are none returns an empty string. Otherwise it returns a test class named from the class and function plus a Branches suffix. The script scripts.serialize_codebase_ast also reaches this code through the testgen package; tests/scripts/test_ast_branches.py covers that use. Other tests are tests/test_orchestrata_tools_testgen_branch_coverage.py, ..._param_strategy.py and ..._units.py.

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/branch.py

"""branch.py — Generate real, runnable branch/exception tests from AST structure.

For every branch event inside a function we emit a complete pytest test —
not a stub — by deriving concrete inputs directly from the condition AST:

  If node   → two tests: one that takes the true path, one the false path.
              Inputs are derived by analysing the condition expression.

  Raise node → one test: inputs that reach the raise + pytest.raises(ExcType).

  Try node   → one test per handler: inputs that trigger that exception class.
              If the handler catches and re-raises a different one, we assert
              on the outer exception.

  Assert node → two tests:
                  (a) valid input → assert does NOT fire
                  (b) violating input → pytest.raises(AssertionError)

Each test is self-contained: it constructs the object (if a method),
calls the function with the derived inputs, and asserts the expected
outcome — no pytest.skip, no pass, no TODO in the body.

Where we cannot fully resolve an input (unknown custom type) we emit
``cast(Any, None)`` with the import registered on the TestFile — never a
mock (see ``types.sentinel``).

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 parser's ``StatementRecord`` tree walk is replaced by a direct
walk of the already-parsed ``ast.FunctionDef``/``ast.AsyncFunctionDef``
node (see ``orchestrata_tools.test_generator.testgen.adapter`` for how
that node is located). ``conditions.py``'s condition analysis is untouched
— it always took a condition string + a plain args dict, never the
deleted parser's types.

Template-defect fixes applied here:
  (a) THE fix. The original (pre-deletion) generator never emitted
      ``async def``/``@pytest.mark.asyncio`` at all — every scaffolded
      test was a plain synchronous ``def``. A hand-patched generated file
      (``tests/generated_recovery/test_agents_volks_qa_volk.py``) shows
      what happens when someone adds the async wrapper by hand without
      the ``await``: ``self._obj.process(...)`` returns a (never awaited)
      coroutine object, so ``assert result is not None`` is vacuously
      true — false confidence, plus a live ``RuntimeWarning``. This
      version makes async-awareness a first-class part of the generator:
      when the target is ``async def``, the test wrapper is emitted as
      ``async def`` with ``@pytest.mark.asyncio``, and the call is
      awaited.
  (c) When the target class's bases resolve, under TYPE_CHECKING, to a
      package-local ``_XHost``-style Protocol with an unimplemented stub
      (the "mixin self-type" trick — see ``adapter.looks_abstract_via_mixin``),
      the constructor call carries ``# pyright: ignore[reportAbstractUsage]``.
"""

from __future__ import annotations

import ast
import re
from typing import TYPE_CHECKING, Any

from orchestrata_tools.test_generator.testgen.conditions import analyse
from orchestrata_tools.test_generator.testgen.types import (
    fixture_args,
    sentinel_for_param,
    sentinel_imports,
    sentinel_needs_cast,
)
from orchestrata_tools.test_generator.testgen.writer import CodeBuilder, TestFile

if TYPE_CHECKING:
    from orchestrata_tools.test_generator.testgen.config import ProjectConfig


def gen_branch_tests(
    func_name: str,
    func_qname: str,
    func_node: ast.FunctionDef | ast.AsyncFunctionDef,
    source: str,
    is_method: bool = False,
    class_name: str | None = None,
    class_node: ast.ClassDef | None = None,
    tf: TestFile | None = None,
    config: ProjectConfig | None = None,
    looks_abstract_via_mixin: bool = False,
) -> str:
    args_info = _args_info(func_node.args)
    returns = ast.unparse(func_node.returns) if func_node.returns else None
    is_async = isinstance(func_node, ast.AsyncFunctionDef)

    events = _collect_events(func_node.body, depth=0)
    if not events:
        return ""

    label = (
        f"Test{class_name}{func_name.title().replace('_', '')}"
        if class_name
        else f"Test{func_name.title().replace('_', '')}"
    ) + "Branches"

    cb = CodeBuilder()
    cb.cls(label)
    # Sanitize source to avoid triple-quote breakage in docstrings
    safe_source = _sanitize_docstring_source(source[:500] if source else "")
    cb.docstring(
        f"Branch/exception tests for: {func_qname}\n"
        f"Every test exercises a specific control-flow path.\n"
        + (f"\nSource:\n{safe_source}" if safe_source else "")
    )
    cb.blank()

    if is_method and class_name:
        cb.func("setup_method")
        ctor = _build_constructor(class_name, class_node, tf, config)
        cb.line(f"# TODO: instantiate {class_name} with valid constructor args")
        ignore = (
            "  # pyright: ignore[reportAbstractUsage]"
            if looks_abstract_via_mixin
            else ""
        )
        cb.line(f"self._obj = {ctor}{ignore}")
        cb.end()
        cb.blank()

    generated: set[str] = set()

    for ev in events:
        for t in _event_to_tests(
            ev, func_name, args_info, returns, is_method, is_async, tf, config
        ):
            if t["name"] in generated:
                continue
            generated.add(t["name"])
            if is_async:
                cb.line("@pytest.mark.asyncio")
            fixtures = fixture_args(t["body"])
            cb.func(
                t["name"],
                args=", ".join(["self", *fixtures]),
                is_async=is_async,
            )
            cb.docstring(t["doc"])
            for ln in t["body"]:
                cb.line(ln)
            cb.end()
            cb.blank()

    if not generated:
        return ""

    cb.end()
    return cb.render()


# ---------------------------------------------------------------------------
# Argument-shape helper (mirrors adapter.args_info_from_ast — kept local so
# branch.py has no import-time dependency on adapter.py's ast.ClassDef
# lookups, only on the plain ast.arguments shape it already receives).
# ---------------------------------------------------------------------------


def _args_info(args: ast.arguments) -> dict[str, Any]:
    def _param(a: ast.arg) -> dict[str, str | None]:
        return {
            "name": a.arg,
            "annotation": ast.unparse(a.annotation) if a.annotation else None,
        }

    return {
        "posonlyargs": [_param(a) for a in args.posonlyargs],
        "args": [_param(a) for a in args.args],
        "kwonlyargs": [_param(a) for a in args.kwonlyargs],
        "defaults_count": len(args.defaults),
        "kw_defaults_count": sum(1 for d in args.kw_defaults if d is not None),
    }


# ---------------------------------------------------------------------------
# Constructor inference
# ---------------------------------------------------------------------------


def _build_constructor(
    class_name: str,
    class_node: ast.ClassDef | None,
    tf: TestFile | None,
    config: ProjectConfig | None,
) -> str:
    """
    Build a constructor call for *class_name* by looking up __init__
    in the class's already-parsed AST body and filling required args
    with sentinels. Falls back to ``ClassName()`` if no __init__ is found.
    """
    if class_node is None:
        return f"{class_name}()"

    init_node = _find_init(class_node)
    if init_node is None:
        return f"{class_name}()"

    args_info = _args_info(init_node.args)
    parts: list[str] = []
    all_args = args_info.get("args", [])
    defaults_count = args_info.get("defaults_count", 0)
    # Required args = all args except the last defaults_count (self included
    # in the count base, but self never has a default and is always index 0)
    required_count = len(all_args) - defaults_count
    for i, a in enumerate(all_args):
        if a["name"] in ("self", "cls"):
            continue
        if i >= required_count:
            continue  # optional argument — skip to keep test minimal; no
            # sentinel computed, so no unused import gets registered for
            # a value that never reaches the generated call (fix (b)'s
            # tracking must stay precise, not just present)
        val = _sentinel_tracked(a.get("annotation"), tf, config)
        parts.append(f"{a['name']}={val}")

    kwonly = args_info.get("kwonlyargs", [])
    kw_defaults_count = args_info.get("kw_defaults_count", 0)
    kw_required_count = len(kwonly) - kw_defaults_count
    for i, a in enumerate(kwonly):
        if i >= kw_required_count:
            continue
        val = _sentinel_tracked(a.get("annotation"), tf, config)
        parts.append(f"{a['name']}={val}")

    if parts:
        return f"{class_name}({', '.join(parts)})"
    return f"{class_name}()"


def _find_init(
    class_node: ast.ClassDef,
) -> ast.FunctionDef | ast.AsyncFunctionDef | None:
    """The class's own ``__init__`` (direct body — a base class's
    ``__init__`` is out of scope here; ``ClassName()`` is the safe
    fallback when this class declares none)."""
    for item in class_node.body:
        if (
            isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef))
            and item.name == "__init__"
        ):
            return item
    return None


# ---------------------------------------------------------------------------
# Event collection — direct ast walk (replaces the deleted parser's
# StatementRecord tree walk one-for-one: same event kinds, same recursion
# rules — stop at nested def/class, otherwise descend into every compound
# statement's body/orelse/handlers/finalbody).
# ---------------------------------------------------------------------------


def _collect_events(
    stmts: list[ast.stmt], depth: int, max_depth: int = 5
) -> list[dict[str, Any]]:
    if depth > max_depth:
        return []
    events: list[dict[str, Any]] = []
    for s in stmts:
        if isinstance(s, ast.If):
            events.append(
                {"kind": "if", "lineno": s.lineno, "test": _try_unparse(s.test)}
            )
            events.extend(_collect_events(s.body, depth + 1, max_depth))
            events.extend(_collect_events(s.orelse, depth + 1, max_depth))
        elif isinstance(s, ast.Try):
            handlers = [
                {
                    "type": _try_unparse(h.type) if h.type else None,
                    "reraise": _handler_reraises(h.body, h.type),
                }
                for h in s.handlers
            ]
            events.append(
                {
                    "kind": "try",
                    "lineno": s.lineno,
                    "handlers": handlers,
                    "has_finally": bool(s.finalbody),
                }
            )
            events.extend(_collect_events(s.body, depth + 1, max_depth))
            for h in s.handlers:
                events.extend(_collect_events(h.body, depth + 1, max_depth))
            events.extend(_collect_events(s.orelse, depth + 1, max_depth))
            events.extend(_collect_events(s.finalbody, depth + 1, max_depth))
        elif isinstance(s, ast.Assert):
            events.append(
                {
                    "kind": "assert",
                    "lineno": s.lineno,
                    "test": _try_unparse(s.test),
                    "msg_node": s.msg,
                }
            )
        elif isinstance(s, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            continue  # nested scope — not a branch of the enclosing function
        elif isinstance(s, ast.Raise):
            events.append(
                {
                    "kind": "raise",
                    "lineno": s.lineno,
                    "exc": _try_unparse(s.exc) if s.exc else None,
                    "exc_node": s.exc,
                }
            )
        else:
            for field_name in ("body", "orelse", "finalbody"):
                child = getattr(s, field_name, None)
                if child:
                    events.extend(_collect_events(child, depth + 1, max_depth))
    return sorted(events, key=lambda e: e["lineno"])


def _try_unparse(node: ast.AST) -> str:
    try:
        return ast.unparse(node)
    except Exception:
        return "?"


def _handler_reraises(
    handler_body: list[ast.stmt], caught_type: ast.expr | None
) -> str | None:
    caught = _exc_class(_try_unparse(caught_type)) if caught_type else "Exception"
    for s in handler_body:
        if isinstance(s, ast.Raise) and s.exc is not None:
            exc = _exc_class(_try_unparse(s.exc))
            if exc != caught:
                return exc
    return None


# ---------------------------------------------------------------------------
# Event → concrete tests
# ---------------------------------------------------------------------------


def _event_to_tests(
    ev: dict[str, Any],
    func_name: str,
    args_info: dict[str, Any],
    returns: str | None,
    is_method: bool,
    is_async: bool,
    tf: TestFile | None = None,
    config: ProjectConfig | None = None,
) -> list[dict[str, Any]]:
    kind = ev["kind"]
    lineno = ev["lineno"]
    tests: list[dict[str, Any]] = []

    call = _build_call(func_name, args_info, is_method, is_async, tf=tf, config=config)

    if kind == "if":
        cond = ev.get("test") or "condition"
        ci = analyse(cond, args_info)

        true_call = _build_call(
            func_name,
            args_info,
            is_method,
            is_async,
            ci.true_inputs,
            tf=tf,
            config=config,
        )
        true_body = list(ci.setup_code) + [true_call] + _assert_return(returns)
        tests.append(
            {
                "name": f"test_L{lineno}_when_{_slug(cond)}_true",
                "doc": f"Line {lineno}: condition `{cond}` is True.\n"
                f"Inputs: {ci.true_inputs or '(see setup code)'}",
                "body": true_body,
            }
        )

        false_call = _build_call(
            func_name,
            args_info,
            is_method,
            is_async,
            ci.false_inputs,
            tf=tf,
            config=config,
        )
        false_body = list(ci.setup_code) + [false_call] + _assert_return(returns)
        tests.append(
            {
                "name": f"test_L{lineno}_when_{_slug(cond)}_false",
                "doc": f"Line {lineno}: condition `{cond}` is False.\n"
                f"Inputs: {ci.false_inputs or '(see setup code)'}",
                "body": false_body,
            }
        )

    elif kind == "raise":
        exc = ev.get("exc") or "Exception"
        exc_class = _exc_class(exc)
        msg_match = _exc_message(ev.get("exc_node"))

        match_part = f", match={msg_match}" if msg_match else ""
        tests.append(
            {
                "name": f"test_raises_{exc_class}_L{lineno}",
                "doc": f"Line {lineno}: {exc} must be raised when its guard condition is met.",
                "body": [
                    f"with pytest.raises({exc_class}{match_part}):",
                    f"    {call}  # provide inputs that reach the raise at line {lineno}",
                ],
            }
        )

    elif kind == "try":
        for h in ev.get("handlers", []):
            exc_type = h.get("type") or "Exception"
            exc_class = _exc_class(exc_type)
            reraises = h.get("reraise")

            if reraises:
                body = [
                    f"# Handler catches {exc_type} and re-raises as {reraises}",
                    f"with pytest.raises({reraises}):",
                    f"    {call}  # trigger the {exc_type} path",
                ]
            else:
                body = [
                    f"# Handler suppresses {exc_type} — call must NOT propagate it",
                    f"{call}  # trigger the {exc_type} path",
                    "# assert recovery state here",
                ]
            tests.append(
                {
                    "name": f"test_handles_{exc_class}_L{lineno}",
                    "doc": f"Line {lineno}: {exc_type} handler — "
                    f"{'re-raises as ' + reraises if reraises else 'suppressed'}.",
                    "body": body,
                }
            )

        if ev.get("has_finally"):
            tests.append(
                {
                    "name": f"test_finally_always_runs_L{lineno}",
                    "doc": f"Line {lineno}: finally block must run on both success and exception paths.",
                    "body": [
                        "# Happy path — finally must still execute",
                        f"{call}",
                        "assert True  # if finally raised, this line would not be reached",
                    ],
                }
            )

    elif kind == "assert":
        cond = ev.get("test") or "condition"
        ci = analyse(cond, args_info)

        # Template-defect (found while restoring): the pre-deletion source
        # passed ci.false_inputs to the "passes" call and ci.true_inputs to
        # the "violation" call — backwards. ci.true_inputs are the inputs
        # that make the CONDITION True, i.e. the assert HOLDS (no error);
        # ci.false_inputs make it False, i.e. the assert RAISES. Swapped,
        # "test_assert_passes" actually drove the violating input and
        # raised AssertionError from inside a test that never wrapped the
        # call in pytest.raises — failing on its own generated output.
        valid_call = _build_call(
            func_name,
            args_info,
            is_method,
            is_async,
            ci.true_inputs,
            tf=tf,
            config=config,
        )
        tests.append(
            {
                "name": f"test_assert_passes_L{lineno}",
                "doc": f"Line {lineno}: assert `{cond}` must hold under valid inputs.",
                "body": list(ci.setup_code)
                + [
                    f"# Inputs where `{cond}` is satisfied",
                    valid_call,
                    "# No AssertionError raised — test passes",
                ],
            }
        )

        bad_call = _build_call(
            func_name,
            args_info,
            is_method,
            is_async,
            ci.false_inputs,
            tf=tf,
            config=config,
        )
        msg_match = _assert_message(ev.get("msg_node"))
        match_part = f", match={msg_match}" if msg_match else ""
        tests.append(
            {
                "name": f"test_assert_violation_L{lineno}",
                "doc": f"Line {lineno}: violating `{cond}` must raise AssertionError.",
                "body": list(ci.setup_code)
                + [
                    f"with pytest.raises(AssertionError{match_part}):",
                    f"    {bad_call}",
                ],
            }
        )

    return tests


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


def _sentinel_tracked(
    annotation: str | None,
    tf: TestFile | None,
    config: ProjectConfig | None = None,
    name: str = "",
) -> str:
    """``types.sentinel`` plus import tracking — the ``cast(Any, None)``
    fallback (template-defect fix (b)) needs ``from typing import Any,
    cast`` registered on the SAME generated file, and any stdlib sentinel
    (``pathlib.Path()``, ``datetime.datetime.now()``, ...) needs its
    module imported too.
    """
    # WI-TESTGEN (task 575): a path-like ``str`` param is built under the
    # ``tmp_path`` fixture, never the bare literal ``"test_value"``.
    val = sentinel_for_param(name, annotation, config)
    if tf is not None:
        for mod in sentinel_imports(annotation):
            tf.add_import(mod)
        if sentinel_needs_cast(annotation, config):
            tf.add_typing_import("Any")
            tf.add_typing_import("cast")
    return val


def _build_call(
    func_name: str,
    args_info: dict[str, Any],
    is_method: bool,
    is_async: bool,
    overrides: dict[str, str] | None = None,
    tf: TestFile | None = None,
    config: ProjectConfig | None = None,
) -> str:
    overrides = overrides or {}
    parts = []
    for a in args_info.get("args", []):
        if a["name"] in ("self", "cls"):
            continue
        val = overrides.get(a["name"]) or _sentinel_tracked(
            a.get("annotation"), tf, config, name=a["name"]
        )
        parts.append(f"{a['name']}={val}")
    for a in args_info.get("kwonlyargs", []):
        val = overrides.get(a["name"]) or _sentinel_tracked(
            a.get("annotation"), tf, config, name=a["name"]
        )
        parts.append(f"{a['name']}={val}")
    args_str = ", ".join(parts)
    prefix = "result = "
    # Template-defect fix (a): await the call when the target is async —
    # otherwise `result` is an unawaited coroutine object, never None, and
    # every `assert result is not None` below is vacuously true.
    await_kw = "await " if is_async else ""
    if is_method:
        return f"{prefix}{await_kw}self._obj.{func_name}({args_str})"
    return f"{prefix}{await_kw}{func_name}({args_str})"


def _assert_return(returns: str | None) -> list[str]:
    if not returns or returns == "None":
        return ["assert result is None"]
    # An Optional/Union-with-None return (very common in this codebase's
    # "X | None" style) makes a plain None a CORRECT outcome for some
    # inputs — this generator's condition analysis (conditions.py) cannot
    # reliably force which branch actually runs for every real-world
    # guard (attribute/collection-state conditions in particular), so
    # asserting `is not None` here produced a live false failure: a
    # genuinely correct `None` return from an Optional-returning function
    # failing its own generated test. Assert only that the call completed.
    if _is_optional_return(returns):
        return [f"# result may legitimately be None here (return type: {returns})"]
    if returns == "bool":
        return ["assert isinstance(result, bool)"]
    if returns == "str":
        return ["assert isinstance(result, str)"]
    if returns == "int":
        return ["assert isinstance(result, int)"]
    if returns.startswith(("list", "List")):
        return ["assert isinstance(result, list)"]
    if returns.startswith(("dict", "Dict")):
        return ["assert isinstance(result, dict)"]
    return [f"assert result is not None  # return type: {returns}"]


def _is_optional_return(returns: str) -> bool:
    """True for ``Optional[X]``/``Union[X, None]`` and the PEP 604
    ``X | None`` spelling this codebase actually writes."""
    r = returns.strip()
    if r.startswith("Optional["):
        return True
    if r.startswith("Union[") and "None" in r:
        return True
    return "None" in [p.strip() for p in r.split("|")]


def _exc_class(exc_expr: str | None) -> str:
    if not exc_expr:
        return "Exception"
    # Handle tuple types like "(FileNotFoundError, PermissionError)"
    expr = exc_expr.strip()
    if expr.startswith("(") and ")" in expr:
        expr = expr[1 : expr.index(")")]
    first = expr.split(",")[0].strip()
    return first.split("(")[0].strip().split(".")[-1]


def _exc_message(exc_node: ast.expr | None) -> str | None:
    """A ``match=`` value for ``pytest.raises``, only when the raised
    exception's first argument is a literal string constant.

    Restoration fix: the pre-deletion version regex-matched the FIRST
    quoted substring anywhere inside the raise's unparsed text — which,
    for something like ``raise Err(f"...", agent_id=d.get('alias', ''))``,
    happily matches a quote pair from an unrelated keyword argument
    (observed live: a bogus ``match="alias', '"`` that made the generated
    ``pytest.raises`` assertion fail on an f-string message it was never
    meant to match). Only a real ``ast.Constant`` string arg is trustworthy
    as a literal match target — an f-string/joined message is content we
    cannot predict without executing the function.
    """
    if not isinstance(exc_node, ast.Call) or not exc_node.args:
        return None
    first = exc_node.args[0]
    if isinstance(first, ast.Constant) and isinstance(first.value, str):
        return repr(re.escape(first.value[:60]))
    return None


def _assert_message(msg_node: ast.expr | None) -> str | None:
    """A ``match=`` value for the ``pytest.raises(AssertionError)`` wrapper
    around a violated ``assert cond, msg`` — same defect class as
    ``_exc_message`` above, found while writing coverage for this exact
    code path. ``ast.unparse(msg_node)`` on a string-constant message
    returns its SOURCE form (quotes included, e.g. ``"'boom'"``), not the
    runtime string the AssertionError actually carries (``"boom"``) — a
    match built from the unparsed form has stray literal quote characters
    the real message never contains, so it never matches and the generated
    test fails on its own output. Only trust a literal ``ast.Constant``
    string, exactly as ``_exc_message`` does for raised exceptions.
    """
    if isinstance(msg_node, ast.Constant) and isinstance(msg_node.value, str):
        return repr(re.escape(msg_node.value[:60]))
    return None


def _slug(condition: str) -> str:
    return re.sub(r"[^a-zA-Z0-9]+", "_", condition).strip("_").lower()[:40] or "cond"


def _sanitize_docstring_source(source: str) -> str:
    """
    Escape triple quotes in source code so it can be safely embedded
    inside a Python docstring.
    """
    if not source:
        return ""
    # Replace " with ' to avoid breaking """ ... """ docstrings
    return source.replace('"""', "'''").replace('"', "'")