orchestrata_tools.test_generator.jinja_gen

Reference › orchestrata_tools › test_generator › jinja_gen

JinjaTestGenerator is a template-driven test generator: deterministic and configurable, with the AST used for @traces extraction, branch detection and method introspection and Jinja2 used to render the output. It supports branch-aware generation for full path coverage. It is re-exported by the package, and its tests are tests/test_orchestrata_tools_jinja_gen.py, tests/test_enum_test_gen_value_agnostic.py and tests/test_orchestrata_tools_testgen_param_strategy.py. The non-template equivalent is enhanced_generator.

generate_module takes classes, the source module and the test module and returns a complete module. _build_context assembles the Jinja context from the classes and a traces map; _method_contexts adds per-method branch data. Branches come from _extract_branches, which reads the cached source file (_get_source), finds the method node (_find_method_node) and walks it (_walk_branches), recording BranchInfo entries with a name, condition and branch type; _format_branches converts them for the template. set_template replaces the default template at run time. Parameter values come from the shared strategy in testgen.types, with a table that also maps NoneType to None. JinjaGenConfig holds the settings.

At a glance

Class diagram

Dependency diagram

API

Classes

BranchInfo (dataclass)

Decorators: dataclass · Line 37

A code branch found in a method body.

JinjaGenConfig (dataclass)

Decorators: dataclass · Line 164

Configuration for Jinja2-based test generation.

JinjaTestGenerator

Line 178

Generate pytest modules using Jinja2 templates + AST introspection.

Source

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

"""Jinja2-powered test generator — deterministic, configurable, template-driven.

Uses AST for @traces extraction, branch detection, and method introspection.
Jinja2 for output rendering. Supports branch-aware test generation for full
path coverage.

Usage:
    from orchestrata_tools.test_generator.jinja_gen import JinjaTestGenerator
    gen = JinjaTestGenerator()
    code = gen.generate_module(classes, "my.module", "test_my_module")
"""

from __future__ import annotations

import ast as py_ast
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING

from jinja2 import BaseLoader, Environment

from orchestrata_tools.test_generator.testgen.types import (
    JINJA_LITERAL_FIXTURES,
    ParamValueStrategy,
    default_param_strategy,
    fixture_args,
)

if TYPE_CHECKING:
    from orchestrata_tools.models.ast_models import ClassInfo, MethodInfo, ParamInfo


# ── Branch info ───────────────────────────────────────────────────────


@dataclass
class BranchInfo:
    """A code branch found in a method body."""

    name: str  # e.g. "when_path_exists", "when_path_none"
    condition: str  # The condition expression as text
    branch_type: str  # "if_true", "if_false", "except", "else", "loop_empty"
    params: list[dict] = field(default_factory=list)  # Parameter overrides
    expected: str = ""  # Expected behavior description


# ── Default Jinja2 template ───────────────────────────────────────────

DEFAULT_TEST_TEMPLATE = '''\
"""Tests for {{ source_module }}.

Test Standards (AGENTS.md § Test Standards):
  T1-T9 enforced. No mocking, stubs, fakes, skips, monkeypatch.
  Function-scope fixtures only. test_id for DB ops. Parallel-safe.
"""

from __future__ import annotations

import pytest
from datetime import UTC, datetime
from pathlib import Path
from uuid import uuid4
{% if import_block %}
{{ import_block }}
{% endif %}
{% if has_traces %}
from orchestrata_ultimate.core.base.traces import traces
{% endif %}

{% for cls in classes %}
{% if cls.traces_deco %}
{{ cls.traces_deco }}
{% endif %}
class Test{{ cls.name }}:
{% if cls.is_enum %}
{% for var in cls.vars %}
    def test_{{ var.name }}_is_valid(self) -> None:
        """Test {{ cls.name }}.{{ var.name }} is a valid enum value."""
        assert {{ cls.name }}.{{ var.name }} in {{ cls.name }}
        assert {{ cls.name }}.{{ var.name }}.name == "{{ var.name }}"

{% endfor %}
{% else %}
    @pytest.fixture
    def obj(self) -> {{ cls.name }}:
        return {{ cls.name }}()

{% for method in cls.methods %}
{# ── base test ── #}
{% if method.traces_deco %}
    {{ method.traces_deco }}
{% endif %}
{% if method.is_async %}
    @pytest.mark.asyncio
    async def test_{{ method.name }}(self, obj: {{ cls.name }}{{ method.fixture_args }}) -> None:
{% else %}
    def test_{{ method.name }}(self, obj: {{ cls.name }}{{ method.fixture_args }}) -> None:
{% endif %}
        """Test {{ cls.name }}.{{ method.name }}()."""

        # Arrange
{% for p in method.params %}
        {{ p.name }} = {{ p.fixture }}
{% endfor %}

        # Act
{% if method.has_return %}
        result = {% if method.is_async %}await {% endif %}obj.{{ method.name }}({{ method.call_args }})
{% else %}
        {% if method.is_async %}await {% endif %}obj.{{ method.name }}({{ method.call_args }})
{% endif %}

        # Assert
{% if method.has_return %}
        assert result is not None
{% else %}
        # Verify side effects or state changes
{% endif %}

{# ── branch tests ── #}
{% for branch in method.branches %}
{% if method.is_async %}
    @pytest.mark.asyncio
    async def test_{{ method.name }}_{{ branch.name }}(self, obj: {{ cls.name }}{{ method.fixture_args }}) -> None:
{% else %}
    def test_{{ method.name }}_{{ branch.name }}(self, obj: {{ cls.name }}{{ method.fixture_args }}) -> None:
{% endif %}
        """Test {{ cls.name }}.{{ method.name }}() — branch: {{ branch.branch_type }} {{ branch.condition }}."""

        # Arrange
{% for p in branch.params %}
        {{ p.name }} = {{ p.fixture }}
{% endfor %}
{% for p in method.params %}
{% if p.name not in branch.param_names %}
        {{ p.name }} = {{ p.fixture }}
{% endif %}
{% endfor %}

        # Act
{% if method.has_return %}
        result = {% if method.is_async %}await {% endif %}obj.{{ method.name }}({{ method.call_args }})
{% else %}
        {% if method.is_async %}await {% endif %}obj.{{ method.name }}({{ method.call_args }})
{% endif %}

        # Assert — {{ branch.expected }}
{% if method.has_return %}
        assert result is not None
{% else %}
        # {{ branch.expected }}
{% endif %}

{% endfor %}
{% endfor %}
{% endif %}
{% endfor %}
'''

# ── Configuration ─────────────────────────────────────────────────────


@dataclass
class JinjaGenConfig:
    """Configuration for Jinja2-based test generation."""

    template: str | None = None
    include_traces: bool = True
    include_imports: bool = True
    include_branches: bool = True  # Generate branch-specific tests
    aaa_pattern: bool = True
    fixture_scope: str = "function"


# ── Generator ─────────────────────────────────────────────────────────


class JinjaTestGenerator:
    """Generate pytest modules using Jinja2 templates + AST introspection.

    Features:
    - @traces decorator extraction and propagation
    - Branch detection (if/else, try/except, match/case, loops)
    - Async method support
    - Enum value testing
    - AAA pattern
    - Configurable via JinjaGenConfig
    """

    #: One source of truth: the literal table lives in ``testgen.types``.
    TYPE_FIXTURES: dict[str, str] = JINJA_LITERAL_FIXTURES

    def __init__(
        self,
        config: JinjaGenConfig | None = None,
        values: ParamValueStrategy | None = None,
    ):
        self.config = config or JinjaGenConfig()
        # WI-TESTGEN (task 575): path-like ``str`` params -> ``tmp_path``.
        self._values = values or default_param_strategy(
            self.TYPE_FIXTURES, "None  # TODO: provide fixture for {name} ({ann})"
        )
        self._env = Environment(loader=BaseLoader())
        self._env.filters["indent"] = lambda s, _n=0: str(s)  # No-op indent
        self._template = self._env.from_string(
            self.config.template or DEFAULT_TEST_TEMPLATE
        )
        self._source_cache: dict[str, str] = {}

    # ── Public API ─────────────────────────────────────────────────

    def generate_module(
        self, classes: list[ClassInfo], source_module: str, test_module: str
    ) -> str:
        """Generate a complete test module with branch coverage."""
        traces_map = self._extract_traces(classes)
        ctx = self._build_context(classes, source_module, traces_map)
        return self._template.render(**ctx)

    def set_template(self, template_str: str) -> None:
        """Override the default template at runtime."""
        self._template = self._env.from_string(template_str)

    # ── Context building ──────────────────────────────────────────

    def _build_context(
        self, classes: list[ClassInfo], source_module: str, traces_map: dict[str, dict]
    ) -> dict:
        """Build Jinja2 template context."""
        has_traces = bool(traces_map)
        import_block = self._build_import_block(classes, source_module)

        cls_ctx = []
        for cls in classes:
            if cls.name.startswith("_"):
                continue
            is_enum = cls.is_enum or (
                not cls.methods
                and cls.class_vars
                and all(
                    hasattr(v, "name") if hasattr(v, "name") else False
                    for v in (cls.class_vars or [])
                )
            )
            class_traces = traces_map.get(cls.name, {})
            cls_ctx.append(
                {
                    "name": cls.name,
                    "is_enum": is_enum,
                    "traces_deco": self._format_traces(class_traces),
                    "vars": self._class_vars(cls),
                    "methods": self._method_contexts(cls, traces_map),
                }
            )

        return {
            "source_module": source_module,
            "has_traces": has_traces,
            "import_block": import_block,
            "classes": cls_ctx,
        }

    def _method_contexts(
        self, cls: ClassInfo, traces_map: dict[str, dict]
    ) -> list[dict]:
        """Build method contexts including branch information."""
        methods = []
        for m in cls.methods:
            if m.name.startswith("_") and m.name != "__init__":
                continue
            params = [p for p in m.parameters if p.name not in ("self", "cls")]
            method_traces = traces_map.get(m.name, {})
            method_params = [
                {"name": p.name, "fixture": self._fixture(p)} for p in params[:5]
            ]
            call_args = ", ".join(f"{p.name}={p.name}" for p in params)

            # Extract branches if enabled
            branches: list[dict] = []
            if self.config.include_branches:
                branches = self._extract_branches(cls, m)

            methods.append(
                {
                    "name": m.name,
                    "is_async": m.is_async,
                    "has_return": bool(m.return_type and m.return_type != "None"),
                    "traces_deco": self._format_traces(method_traces),
                    "params": method_params,
                    "fixture_args": "".join(
                        f", {name}"
                        for name in fixture_args(p["fixture"] for p in method_params)
                    ),
                    "call_args": call_args,
                    "branches": branches,
                }
            )
        return methods

    # ── Branch extraction ─────────────────────────────────────────

    def _extract_branches(self, cls: ClassInfo, method: MethodInfo) -> list[dict]:
        """Extract branch points from a method's AST body.

        Walks the method's AST looking for if/else, try/except, match/case,
        and for/while structures. For each branch, generates parameter
        overrides that would exercise that path.
        """
        source = self._get_source(
            getattr(cls, "file_path", "") if hasattr(cls, "file_path") else ""
        )
        if not source:
            return []

        try:
            tree = py_ast.parse(source)
        except Exception:
            return []

        # Find the method node
        method_node = self._find_method_node(tree, cls.name, method.name)
        if not method_node:
            return []

        branches: list[BranchInfo] = []
        self._walk_branches(method_node, branches, prefix=method.name)
        return self._format_branches(branches)

    def _format_branches(self, branches: list[BranchInfo]) -> list[dict]:
        """Convert BranchInfo to template-ready dicts."""
        return [
            {
                "name": b.name,
                "condition": b.condition[:60],
                "branch_type": b.branch_type,
                "expected": b.expected or "Verify branch behavior",
                "params": b.params,
                "param_names": {p["name"] for p in b.params} if b.params else set(),
            }
            for b in branches
        ]

    def _get_source(self, file_path: str) -> str:
        """Get cached source for a file."""
        if file_path not in self._source_cache:
            try:
                self._source_cache[file_path] = Path(file_path).read_text()
            except Exception:
                return ""
        return self._source_cache[file_path]

    def _find_method_node(
        self, tree: py_ast.AST, class_name: str, method_name: str
    ) -> py_ast.AST | None:
        """Find a method's AST node within its class."""
        for node in py_ast.walk(tree):
            if isinstance(node, py_ast.ClassDef) and node.name == class_name:
                for item in node.body:
                    if isinstance(item, (py_ast.FunctionDef, py_ast.AsyncFunctionDef)):
                        if item.name == method_name:
                            return item
        return None

    def _walk_branches(
        self, node: py_ast.AST, branches: list[BranchInfo], prefix: str
    ) -> None:
        """Walk an AST node and its children to find branch points."""
        # If statement
        if isinstance(node, py_ast.If):
            cond_text = (
                py_ast.unparse(node.test) if hasattr(node, "test") else "condition"
            )
            # True branch
            branches.append(
                BranchInfo(
                    name=f"{prefix}_when_{self._slug(cond_text)}_true",
                    condition=cond_text,
                    branch_type="if_true",
                    expected=f"Condition '{cond_text}' is truthy",
                )
            )
            # False branch (elif or else)
            if node.orelse:
                branches.append(
                    BranchInfo(
                        name=f"{prefix}_when_{self._slug(cond_text)}_false",
                        condition=f"not ({cond_text})",
                        branch_type="if_false",
                        expected=f"Condition '{cond_text}' is falsy",
                    )
                )
            # Walk children (declared as AST so the `child` name can be
            # reused below for py_ast.walk()/iter_child_nodes(), which
            # yield the broader AST type rather than the stmt subtype).
            child: py_ast.AST
            for child in [*node.body, *node.orelse]:
                self._walk_branches(child, branches, prefix)

        # Try/except
        elif isinstance(node, py_ast.Try):
            branches.append(
                BranchInfo(
                    name=f"{prefix}_try_succeeds",
                    condition="no exception",
                    branch_type="try",
                    expected="No exception is raised",
                )
            )
            for handler in node.handlers:
                exc_name = py_ast.unparse(handler.type) if handler.type else "Exception"
                branches.append(
                    BranchInfo(
                        name=f"{prefix}_raises_{self._slug(exc_name)}",
                        condition=exc_name,
                        branch_type="except",
                        params=[
                            {"name": "self", "fixture": "self"},
                            {"name": "obj", "fixture": "obj"},
                        ],
                        expected=f"Raises {exc_name}",
                    )
                )
            for child in py_ast.walk(node):
                if child is not node:
                    self._walk_branches(child, branches, prefix)

        # For / While loops
        elif isinstance(node, (py_ast.For, py_ast.While)):
            branches.append(
                BranchInfo(
                    name=f"{prefix}_iterates",
                    condition=(
                        "has items"
                        if isinstance(node, py_ast.For)
                        else "condition true"
                    ),
                    branch_type="loop_body",
                    expected="Loop body executes",
                )
            )
            branches.append(
                BranchInfo(
                    name=f"{prefix}_empty",
                    condition=(
                        "no items"
                        if isinstance(node, py_ast.For)
                        else "condition false"
                    ),
                    branch_type="loop_empty",
                    expected="Loop body does not execute",
                )
            )

        # Match/case (Python 3.10+)
        elif isinstance(node, py_ast.Match):
            for i, case in enumerate(node.cases):
                pat = (
                    py_ast.unparse(case.pattern)
                    if hasattr(case, "pattern")
                    else f"case_{i}"
                )
                branches.append(
                    BranchInfo(
                        name=f"{prefix}_matches_{self._slug(pat)}",
                        condition=f"match {pat}",
                        branch_type="match_case",
                        expected=f"Matches pattern: {pat}",
                    )
                )

        # Walk children for nested branches
        for child in py_ast.iter_child_nodes(node):
            self._walk_branches(child, branches, prefix)

    @staticmethod
    def _slug(text: str) -> str:
        """Convert a string to a safe identifier slug."""
        import re

        text = re.sub(r"[^a-zA-Z0-9_\s]", "", text[:40])
        text = re.sub(r"\s+", "_", text.strip())
        return text[:30].lower() or "condition"

    # ── Enum + class var helpers ─────────────────────────────────

    def _class_vars(self, cls: ClassInfo) -> list[dict]:
        vars_list = []
        for v in cls.class_vars or []:
            name = getattr(v, "name", "") if hasattr(v, "name") else str(v)
            if name and not name.startswith("_"):
                vars_list.append({"name": name})
        return vars_list

    # ── Import block ──────────────────────────────────────────────

    def _build_import_block(self, classes: list[ClassInfo], module: str) -> str:
        names = [c.name for c in classes if not c.name.startswith("_")]
        if not names:
            return ""
        if len(names) == 1:
            return f"from orchestrata_ultimate.{module} import {names[0]}"
        lines = [f"from orchestrata_ultimate.{module} import ("]
        for name in names:
            lines.append(f"    {name},")
        lines.append(")")
        return "\n".join(lines)

    # ── Fixtures ──────────────────────────────────────────────────

    def _fixture(self, param: ParamInfo) -> str:
        default = getattr(param, "default", None)
        if default is not None:
            return str(default)
        return self._values.value_for(param)

    # ── @traces formatting ────────────────────────────────────────

    def _format_traces(self, meta: dict) -> str | None:
        if not meta or not self.config.include_traces:
            return None
        return (
            f'@traces(spec="{meta.get("spec", "???")}", '
            f'story="{meta.get("story", "???")}", '
            f'reqs=["{meta.get("reqs", "REQ-???-???")}"])'
        )

    # ── AST @traces extraction ────────────────────────────────────

    def _extract_traces(self, classes: list[ClassInfo]) -> dict[str, dict]:
        result: dict[str, dict] = {}
        seen: set[str] = set()
        for cls in classes:
            fp = getattr(cls, "file_path", "") if hasattr(cls, "file_path") else ""
            if not fp or fp in seen:
                continue
            seen.add(fp)
            try:
                result.update(self._parse_traces_file(fp))
            except Exception:
                pass
        return result

    def _parse_traces_file(self, file_path: str) -> dict[str, dict]:
        result: dict[str, dict] = {}
        try:
            with open(file_path) as f:
                tree = py_ast.parse(f.read())
        except Exception:
            return result
        for node in py_ast.walk(tree):
            if not isinstance(
                node,
                (
                    py_ast.ClassDef,
                    py_ast.FunctionDef,
                    py_ast.AsyncFunctionDef,
                ),
            ):
                continue
            for dec in node.decorator_list:
                meta = self._parse_traces_decorator(dec)
                if meta:
                    result[node.name] = meta
        return result

    def _parse_traces_decorator(self, decorator: py_ast.expr) -> dict | None:
        if (
            not isinstance(decorator, py_ast.Call)
            or not isinstance(decorator.func, py_ast.Name)
            or decorator.func.id != "traces"
        ):
            return None
        meta: dict[str, str] = {}
        for kw in decorator.keywords:
            if kw.arg == "spec" and isinstance(kw.value, py_ast.Constant):
                meta["spec"] = str(kw.value.value)
            elif kw.arg == "story" and isinstance(kw.value, py_ast.Constant):
                meta["story"] = str(kw.value.value)
            elif kw.arg == "reqs" and isinstance(kw.value, py_ast.List):
                reqs = [
                    str(e.value)
                    for e in kw.value.elts
                    if isinstance(e, py_ast.Constant)
                ]
                meta["reqs"] = ", ".join(reqs[:3])
        return meta if meta else None