orchestrata_tools.test_generator.stub_generator

Reference › orchestrata_tools › test_generator › stub_generator

TestStubGenerator writes pytest stubs from method signatures. It is used by combined_generator, re-exported by the package, and covered by tests/test_orchestrata_tools_stub_generator.py and tests/test_orchestrata_tools_testgen_param_strategy.py.

The constructor takes a module prefix (default tests) and an optional ParamValueStrategy from testgen.types; path-like string parameters become pytest tmp_path values. generate_stub returns one stub for a method, skipping self and cls in the parameter list. generate_class_stubs covers the public methods of a class, and generate_module_stubs builds a complete test file for several classes. _fixture_for_type suggests a fixture for an annotation.

At a glance

Class diagram

Dependency diagram

API

Classes

TestStubGenerator

Line 19

Generate pytest test stubs from AST method metadata.

Source

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

"""Generate pytest test stubs from method signatures."""

from __future__ import annotations

from typing import TYPE_CHECKING

from orchestrata_tools.models.ast_models import ParamInfo
from orchestrata_tools.test_generator.testgen.types import (
    STUB_LITERAL_FIXTURES,
    ParamValueStrategy,
    default_param_strategy,
    fixture_args,
)

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


class TestStubGenerator:
    """Generate pytest test stubs from AST method metadata."""

    def __init__(
        self,
        module_prefix: str = "tests",
        values: ParamValueStrategy | None = None,
    ):
        self.module_prefix = module_prefix
        # WI-TESTGEN (task 575): path-like ``str`` params -> ``tmp_path``.
        self._values = values or default_param_strategy(
            STUB_LITERAL_FIXTURES, "None  # TODO: provide value"
        )

    def _param_value(self, param: ParamInfo) -> str:
        if not param.type_annotation:
            return "None"
        return self._values.value_for(param)

    def _fixture_for_type(self, type_name: str | None) -> str:
        """Suggest a fixture for a given type annotation."""
        if not type_name:
            return "None"
        return self._values.value_for(ParamInfo(name="", type_annotation=type_name))

    def _param_list(self, method: MethodInfo) -> str:
        """Build parameter list for test call, skipping self/cls."""
        args = []
        for p in method.parameters:
            if p.name in ("self", "cls"):
                continue
            args.append(f"{p.name}={self._param_value(p)}")
        return ", ".join(args)

    def generate_stub(self, method: MethodInfo) -> str:
        """Generate a single test stub for a method."""
        test_name = f"test_{method.name}"
        if method.class_name:
            test_name = f"test_{method.class_name.lower()}_{method.name}"

        qualified_name = (
            f"{method.class_name}.{method.name}" if method.class_name else method.name
        )
        fixtures = ", ".join(
            fixture_args(
                self._param_value(p)
                for p in method.parameters
                if p.name not in ("self", "cls")
            )
        )
        lines = [
            f"def {test_name}({fixtures}):",
            f'    """Test {qualified_name}()."""',
            "    # Arrange",
        ]

        if method.class_name:
            lines.append(f"    obj = {method.class_name}()")
            call_expr = f"obj.{method.name}"
        else:
            call_expr = method.name

        params = self._param_list(method)
        lines.append("    # Act")

        if method.is_async:
            lines.append(f"    result = asyncio.run({call_expr}({params}))")
        else:
            lines.append(f"    result = {call_expr}({params})")

        lines.extend(
            [
                "    # Assert",
                "    assert result is not None  # TODO: assert actual behavior",
                "",
            ]
        )
        return "\n".join(lines)

    def generate_class_stubs(self, cls: ClassInfo) -> str:
        """Generate test stubs for all public methods in a class."""
        lines = [
            f"class Test{cls.name}:",
            f'    """Tests for {cls.name}."""',
            "",
        ]
        for method in cls.methods:
            if not method.name.startswith("_"):
                stub_lines = self.generate_stub(method).split("\n")
                # Indent the stub
                for line in stub_lines:
                    lines.append(f"    {line}")
                lines.append("")
        return "\n".join(lines)

    def generate_module_stubs(self, classes: list[ClassInfo]) -> str:
        """Generate a complete test file with stubs for multiple classes."""
        lines = [
            '"""Generated test stubs."""',
            "",
            "import pytest",
            "import asyncio",
            "import uuid",
            "from datetime import datetime",
            "",
        ]
        for cls in classes:
            lines.append(self.generate_class_stubs(cls))
            lines.append("")
        return "\n".join(lines)