"""ASV benchmarks for tool-add persistence (issue #621). Measures the round-trip: create a tool ➜ list from a fresh session. """ from __future__ import annotations import importlib import os import sys import tempfile from datetime import datetime from pathlib import Path from typing import Any # Ensure the local *source* tree is importable even when ASV has an # older build of the package installed. _SRC = str(Path(__file__).resolve().parents[1] / "src") if _SRC not in sys.path: sys.path.insert(0, _SRC) # Force-reload so ASV picks up the source tree version. import cleveragents # noqa: E402 importlib.reload(cleveragents) from sqlalchemy import create_engine # noqa: E402 from sqlalchemy.engine import Engine # noqa: E402 from sqlalchemy.orm import Session, sessionmaker # noqa: E402 from cleveragents.infrastructure.database.models import Base # noqa: E402 from cleveragents.infrastructure.database.repositories import ( # noqa: E402 ToolRegistryRepository, ) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_tool(name: str) -> dict[str, Any]: now_iso: str = datetime.now().isoformat() ns: str = name.split("/", 1)[0] if "/" in name else "" short: str = name.split("/", 1)[1] if "/" in name else name return { "name": name, "namespace": ns, "short_name": short, "description": f"Bench tool {name}", "tool_type": "tool", "source": "custom", "timeout": 300, "created_at": now_iso, "updated_at": now_iso, "resource_bindings": [], } # --------------------------------------------------------------------------- # Benchmark suite # --------------------------------------------------------------------------- class ToolAddPersistSuite: """Benchmark the add-then-list persistence round-trip.""" timeout = 60 def setup(self) -> None: fd, self._db_path = tempfile.mkstemp(suffix=".db", prefix="tool_bench_") os.close(fd) engine: Engine = create_engine(f"sqlite:///{self._db_path}", echo=False) Base.metadata.create_all(engine) engine.dispose() def teardown(self) -> None: if os.path.exists(self._db_path): os.unlink(self._db_path) def _repo(self) -> ToolRegistryRepository: engine: Engine = create_engine(f"sqlite:///{self._db_path}", echo=False) factory: sessionmaker[Session] = sessionmaker( bind=engine, expire_on_commit=False ) return ToolRegistryRepository(session_factory=factory) # -- tracking benchmark ------------------------------------------------ def track_list_after_add_count(self) -> int: """Returns the count of tools after a single add (should be 1).""" repo: ToolRegistryRepository = self._repo() repo.create(_make_tool("bench/persist-check")) # Open a new repo (fresh engine) to prove persistence repo2: ToolRegistryRepository = self._repo() tools: list[Any] = repo2.list_all() return len(tools) track_list_after_add_count.unit = "tools" # type: ignore[attr-defined]