orchestrata_tools.test_generator.enhanced_generator

Reference › orchestrata_tools › test_generator › enhanced_generator

EnhancedTestGenerator produces spec-compliant pytest modules from class metadata. Its docstring ties it to the project testing standards (SPEC-097 and the test rules in AGENTS.md): no mocking or skips, @traces decorators propagated from the source, function-scope fixtures, a test_id fixture for database operations, pytest.mark.asyncio for async tests, and the arrange-act-assert pattern. Importers of the module are the package and its tests (tests/test_orchestrata_tools_test_generator.py, tests/test_enum_test_gen_value_agnostic.py, tests/test_orchestrata_tools_testgen_param_strategy.py).

TestGenConfig holds the settings. generate_module takes classes, the source module name and the test module name and returns a full test module: a docstring, an ordered import block (standard library, third party, local) from _imports, and one test class per source class via _test_class and _test_method. Enum classes are handled separately. _collect_traces, _parse_traces and _extract_traces read @traces decorators from the source file and _traces_decorator re-emits them on the generated tests. Parameter values come from a ParamValueStrategy in testgen.types, by default the table LITERAL_FIXTURES. A template-based alternative is jinja_gen.

At a glance

Class diagram

Dependency diagram

API

Classes

TestGenConfig (dataclass)

Decorators: dataclass · Line 30

Configuration for test generation.

EnhancedTestGenerator

Line 43

Generate spec-compliant pytest stubs with full @traces support.

Source

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

"""Enhanced test generator — produces spec-compliant pytest stubs with @traces, async support, fixtures, and AAA pattern.

Follows project standards: SPEC-097 (Testing), AGENTS.md § Test Standards (T1-T14).
- No mocking, stubs, fakes, skips
- @traces decorators propagated from source
- Function-scope fixtures
- test_id fixture for DB ops
- Async tests use @pytest.mark.asyncio
- AAA pattern (Arrange-Act-Assert)
"""

from __future__ import annotations

import ast as py_ast
from dataclasses import dataclass
from typing import TYPE_CHECKING

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

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


@dataclass
class TestGenConfig:
    """Configuration for test generation."""

    include_traces: bool = True
    include_imports: bool = True
    include_fixtures: bool = True
    async_support: bool = True
    fixture_scope: str = "function"
    aaa_pattern: bool = True
    module_docstring: bool = True
    test_class_prefix: str = "Test"


class EnhancedTestGenerator:
    """Generate spec-compliant pytest stubs with full @traces support."""

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

    def __init__(
        self,
        config: TestGenConfig | None = None,
        values: ParamValueStrategy | None = None,
    ):
        self.config = config or TestGenConfig()
        # WI-TESTGEN (task 575): path-like ``str`` params are bound to
        # ``tmp_path`` so a generated test can never write into the checkout.
        self._values = values or default_param_strategy(
            self.TYPE_FIXTURE, "None  # TODO: {name} ({ann})"
        )

    # ── Module-level generation ────────────────────────────────────

    def generate_module(
        self, classes: list[ClassInfo], source_module: str, test_module: str
    ) -> str:
        """Generate a complete test module."""
        parts: list[str] = []

        if self.config.module_docstring:
            parts.append(self._docstring(source_module))

        if self.config.include_imports:
            parts.append(self._imports(classes, source_module))

        all_traces = self._collect_traces(classes)

        for cls_info in classes:
            if cls_info.name.startswith("_"):
                continue
            is_enum = cls_info.is_enum
            parts.append(
                self._test_class(cls_info, source_module, all_traces, is_enum=is_enum)
            )

        return "\n\n".join(parts) + "\n"

    def _docstring(self, source_module: str) -> str:
        return (
            f'"""Tests for {source_module}.\n\n'
            "Test Standards (AGENTS.md § Test Standards):\n"
            "  T1-T9 enforced. No mocking, stubs, fakes, skips, monkeypatch.\n"
            "  Function-scope fixtures only. test_id for DB ops. Parallel-safe.\n"
            '"""\n'
        )

    def _imports(self, classes: list[ClassInfo], source_module: str) -> str:
        """Generate ordered import block: stdlib → third-party → local."""
        lines = []
        lines.append("from __future__ import annotations")
        lines.append("")
        lines.append("import pytest")
        lines.append("from datetime import UTC, datetime")
        lines.append("from pathlib import Path")
        lines.append("from uuid import uuid4")

        # Source module imports
        cls_names = [c.name for c in classes if not c.name.startswith("_")]
        if len(cls_names) == 1:
            lines.append(
                f"from orchestrata_ultimate.{source_module} import {cls_names[0]}"
            )
        else:
            lines.append(f"from orchestrata_ultimate.{source_module} import (")
            lines.extend(f"    {name}," for name in cls_names)
            lines.append(")")

        # @traces import if needed
        if self._has_traces(classes):
            lines.append("from orchestrata_ultimate.core.base.traces import traces")

        return "\n".join(lines) + "\n"

    def _has_traces(self, classes: list[ClassInfo]) -> bool:
        for cls in classes:
            fp = getattr(cls, "file_path", "") if hasattr(cls, "file_path") else ""
            if not fp:
                continue
            try:
                with open(fp) as f:
                    if "@traces" in f.read():
                        return True
            except Exception:
                pass
        return False

    # ── Test class ─────────────────────────────────────────────────

    def _test_class(
        self,
        cls_info: ClassInfo,
        source_module: str,
        all_traces: dict,
        is_enum: bool = False,
    ) -> str:
        lines = []
        test_name = f"{self.config.test_class_prefix}{cls_info.name}"

        class_traces = all_traces.get(cls_info.name, {})
        if class_traces and self.config.include_traces:
            lines.append(self._traces_decorator(class_traces, indent_level=0))

        if is_enum:
            lines.append(f"class {test_name}:")
            for var in cls_info.class_vars or []:
                vname = getattr(var, "name", "") if hasattr(var, "name") else str(var)
                lines.append(f"    def test_{vname.lower()}_is_valid(self) -> None:")
                lines.append(
                    f'        """Test {cls_info.name}.{vname} is a valid enum value."""'
                )
                # Value-agnostic: don't assume the member's value is its
                # lowercased name (e.g. RED = "#ff0000", ACTIVE = 1, auto()).
                lines.append(
                    f"        assert {cls_info.name}.{vname} in {cls_info.name}"
                )
                lines.append(
                    f'        assert {cls_info.name}.{vname}.name == "{vname}"'
                )
                lines.append("")
            return "\n".join(lines).rstrip()

        lines.append(f"class {test_name}:")
        lines.append("")
        lines.append("    @pytest.fixture")
        lines.append(f"    def obj(self) -> {cls_info.name}:")
        lines.append(f"        return {cls_info.name}()")
        lines.append("")

        for method in cls_info.methods:
            if method.name.startswith("_") and method.name != "__init__":
                continue
            m = self._test_method(method, cls_info, all_traces)
            if m:
                lines.append(m)

        return "\n".join(lines).rstrip()

    def _test_method(
        self,
        method: MethodInfo,
        cls_info: ClassInfo,
        all_traces: dict,
    ) -> str:
        lines = []
        method_traces = all_traces.get(f"{cls_info.name}.{method.name}", {})
        if method_traces and self.config.include_traces:
            lines.append(self._traces_decorator(method_traces, indent_level=1))

        if self.config.async_support and method.is_async:
            lines.append("    @pytest.mark.asyncio")

        # Cap at 5 params: the Act call below MUST use this same capped
        # list, not the full `params` — otherwise a method with >5
        # params generates a call referencing names past the cap that
        # the Arrange section never assigned (NameError at test-run time).
        params = [p for p in method.parameters if p.name not in ("self", "cls")][:5]
        values = [self._fixture_for(p) for p in params]
        extra = "".join(f", {name}" for name in fixture_args(values))

        test_name = f"test_{method.name}"
        if method.is_async:
            lines.append(
                f"    async def {test_name}(self, obj: {cls_info.name}{extra}) -> None:"
            )
        else:
            lines.append(
                f"    def {test_name}(self, obj: {cls_info.name}{extra}) -> None:"
            )

        lines.append(f'        """Test {cls_info.name}.{method.name}()."""')

        if self.config.aaa_pattern:
            lines.append("")
            lines.append("        # Arrange")
            lines.extend(
                f"        {p.name} = {value}" for p, value in zip(params, values)
            )

            lines.append("")
            lines.append("        # Act")
            call_args = ", ".join(f"{p.name}={p.name}" for p in params)
            async_kw = "await " if method.is_async else ""
            if method.return_type and method.return_type != "None":
                lines.append(
                    f"        result = {async_kw}obj.{method.name}({call_args})"
                )
            else:
                lines.append(f"        {async_kw}obj.{method.name}({call_args})")

            lines.append("")
            lines.append("        # Assert")
            if method.return_type and method.return_type != "None":
                lines.append("        assert result is not None")
            else:
                lines.append("        # Verify side effects or state changes")
        else:
            lines.append("        pass  # TODO: implement")

        return "\n".join(lines) + "\n"

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

    def _traces_decorator(self, meta: dict, indent_level: int = 0) -> str:
        prefix = "    " * indent_level
        reqs = meta.get("reqs") or ["REQ-???-???"]
        # `reqs` is normally a list (see _extract_traces), but tolerate a
        # bare string too rather than iterating its characters.
        if isinstance(reqs, str):
            reqs = [reqs]
        reqs_literal = ", ".join(f'"{r}"' for r in reqs)
        return (
            f"{prefix}@traces("
            f'spec="{meta.get("spec", "???")}", '
            f'story="{meta.get("story", "???")}", '
            f"reqs=[{reqs_literal}])"
        )

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

    def _collect_traces(self, classes: list[ClassInfo]) -> dict[str, dict]:
        result: dict[str, dict] = {}
        for cls in classes:
            fp = getattr(cls, "file_path", "") if hasattr(cls, "file_path") else ""
            if fp:
                try:
                    result.update(self._parse_traces(fp))
                except Exception:
                    pass
        return result

    def _parse_traces(self, file_path: str) -> dict[str, dict]:
        """Map class/method names to their @traces metadata.

        Keyed to match what ``_test_class``/``_test_method`` look up:
        a bare class name ("Widget") for class-level @traces, and
        "ClassName.method_name" for method-level @traces. A flat
        ``ast.walk`` (used previously) loses the enclosing-class context,
        so it keyed everything by bare node name — colliding methods of
        the same name across different classes and never matching the
        qualified-name lookups the callers actually perform (@traces was
        silently never emitted). This walks module body -> class body
        explicitly to preserve that context.
        """
        result: dict[str, dict] = {}
        try:
            with open(file_path) as f:
                tree = py_ast.parse(f.read())
        except Exception:
            return result

        for node in tree.body:
            if isinstance(node, py_ast.ClassDef):
                for dec in node.decorator_list:
                    meta = self._extract_traces(dec)
                    if meta:
                        result[node.name] = meta
                for item in node.body:
                    if isinstance(item, (py_ast.FunctionDef, py_ast.AsyncFunctionDef)):
                        for dec in item.decorator_list:
                            meta = self._extract_traces(dec)
                            if meta:
                                result[f"{node.name}.{item.name}"] = meta
            elif isinstance(node, (py_ast.FunctionDef, py_ast.AsyncFunctionDef)):
                for dec in node.decorator_list:
                    meta = self._extract_traces(dec)
                    if meta:
                        result[node.name] = meta

        return result

    def _extract_traces(
        self, decorator: py_ast.expr
    ) -> dict[str, str | list[str]] | 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 | list[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)
                ]
                # Keep as a real list (capped at 3) — joining into one
                # comma-separated string here used to produce a single
                # malformed reqs=["A, B, C"] entry instead of a proper
                # multi-element list, which the real @traces() decorator
                # (core.base.traces) rejects: it validates each requirement
                # ID individually against a strict REQ-NNN-NNN pattern.
                meta["reqs"] = reqs[:3]
        return meta if meta else None