cleveragents-core/benchmarks/tool_registry_bench.py

"""ASV benchmarks for Tool Registry persistence operations.

Measures the performance of:
- ToolModel.from_domain() construction
- ToolModel.to_domain() reconstruction
- ToolRegistryRepository CRUD operations (in-memory SQLite)
- ValidationAttachmentRepository attach/list operations
"""

from __future__ import annotations

import sys
from datetime import datetime
from pathlib import Path

try:
    from sqlalchemy import create_engine
    from sqlalchemy.orm import sessionmaker

    from cleveragents.infrastructure.database.models import Base, ToolModel
    from cleveragents.infrastructure.database.repositories import (
        ToolRegistryRepository,
        ValidationAttachmentRepository,
    )
except ModuleNotFoundError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
    from sqlalchemy import create_engine
    from sqlalchemy.orm import sessionmaker

    from cleveragents.infrastructure.database.models import Base, ToolModel
    from cleveragents.infrastructure.database.repositories import (
        ToolRegistryRepository,
        ValidationAttachmentRepository,
    )


def _make_tool_dict(
    name: str = "bench/tool-test",
    source: str = "builtin",
    tool_type: str = "tool",
) -> dict[str, object]:
    """Create a minimal tool dict for benchmarks."""
    now = datetime.now().isoformat()
    return {
        "name": name,
        "description": f"Benchmark tool {name}",
        "tool_type": tool_type,
        "source": source,
        "timeout": 300,
        "created_at": now,
        "updated_at": now,
        "resource_bindings": [
            {
                "slot_name": "repo",
                "resource_type": "git-checkout",
                "access_mode": "read_only",
                "binding_mode": "contextual",
                "required": True,
            },
        ],
    }


def _setup_db():
    """Create in-memory DB and return (session, tool_repo, att_repo)."""
    engine = create_engine("sqlite:///:memory:", echo=False)
    Base.metadata.create_all(engine)
    factory = sessionmaker(bind=engine, expire_on_commit=False)
    session = factory()
    tool_repo = ToolRegistryRepository(session_factory=lambda: session)
    att_repo = ValidationAttachmentRepository(session_factory=lambda: session)
    return session, tool_repo, att_repo


class ToolModelSuite:
    """Benchmark ToolModel ORM construction and conversion."""

    def setup(self) -> None:
        """Prepare objects."""
        self.tool_dict = _make_tool_dict()

    def time_from_domain(self) -> None:
        """Benchmark ToolModel.from_domain() construction."""
        ToolModel.from_domain(self.tool_dict)

    def time_to_domain(self) -> None:
        """Benchmark ToolModel.to_domain() conversion."""
        model = ToolModel.from_domain(self.tool_dict)
        model.resource_bindings_rel = []
        model.to_domain()


class ToolRegistryCRUDSuite:
    """Benchmark ToolRegistryRepository CRUD operations."""

    def setup(self) -> None:
        """Prepare a fresh in-memory database for each benchmark."""
        self.session, self.repo, self.att_repo = _setup_db()
        self._counter = 0

    def time_create_tool(self) -> None:
        """Benchmark creating a single tool."""
        self._counter += 1
        tool = _make_tool_dict(name=f"bench/tool-{self._counter}")
        self.repo.create(tool)
        self.session.commit()

    def time_get_by_name(self) -> None:
        """Benchmark retrieving a tool by name."""
        name = "bench/lookup-target"
        if self.repo.get_by_name(name) is None:
            self.repo.create(_make_tool_dict(name=name))
            self.session.commit()
        self.repo.get_by_name(name)

    def time_list_all(self) -> None:
        """Benchmark listing all tools."""
        # Ensure some data exists
        for i in range(5):
            n = f"bench/list-{i}"
            if self.repo.get_by_name(n) is None:
                self.repo.create(_make_tool_dict(name=n))
        self.session.commit()
        self.repo.list_all()


class ValidationAttachmentSuite:
    """Benchmark ValidationAttachmentRepository operations."""

    def setup(self) -> None:
        """Prepare database with a validation tool."""
        self.session, self.repo, self.att_repo = _setup_db()
        val = _make_tool_dict(
            name="bench/val-bench",
            tool_type="validation",
        )
        val["mode"] = "required"
        self.repo.create(val)
        self.session.commit()
        self._counter = 0

    def time_attach(self) -> None:
        """Benchmark attaching a validation."""
        self._counter += 1
        self.att_repo.attach(
            validation_name="bench/val-bench",
            resource_id=f"res-{self._counter}",
            mode="required",
        )
        self.session.commit()

    def time_list_for_resource(self) -> None:
        """Benchmark listing attachments for a resource."""
        res_id = "res-bench-list"
        if not self.att_repo.list_for_resource(res_id):
            self.att_repo.attach(
                validation_name="bench/val-bench",
                resource_id=res_id,
                mode="required",
            )
            self.session.commit()
        self.att_repo.list_for_resource(res_id)