cleveragents-core/robot/helper_tool_add_persist.py

"""Helper script for tool_add_persist.robot integration tests.

Each subcommand exercises the real ToolRegistryRepository against a
file-based SQLite database to verify that data survives session disposal.
"""

from __future__ import annotations

import os
import sys
import tempfile
from datetime import datetime
from pathlib import Path
from typing import Any

# Ensure local source tree is importable
_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

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,
)


def _make_tool(name: str) -> dict[str, Any]:
    """Build a minimal tool dict."""
    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"Integration test tool {name}",
        "tool_type": "tool",
        "source": "custom",
        "timeout": 300,
        "created_at": now_iso,
        "updated_at": now_iso,
        "resource_bindings": [],
    }


def _engine(db_path: str) -> Engine:
    return create_engine(f"sqlite:///{db_path}", echo=False)


def _factory(engine: Engine) -> sessionmaker[Session]:
    return sessionmaker(bind=engine, expire_on_commit=False)


# -------------------------------------------------------------------
# Subcommands
# -------------------------------------------------------------------


def add_then_list() -> None:
    """Add a tool, open a new session, list — should find it."""
    fd, db_path = tempfile.mkstemp(suffix=".db", prefix="tool_persist_robot_")
    os.close(fd)
    try:
        eng: Engine = _engine(db_path)
        Base.metadata.create_all(eng)
        repo1: ToolRegistryRepository = ToolRegistryRepository(
            session_factory=_factory(eng),
        )
        repo1.create(_make_tool("local/robot-test"))
        eng.dispose()

        eng2: Engine = _engine(db_path)
        repo2: ToolRegistryRepository = ToolRegistryRepository(
            session_factory=_factory(eng2),
        )
        tools: list[Any] = repo2.list_all()
        eng2.dispose()

        names: list[str] = [
            str(t.get("name", "") if isinstance(t, dict) else getattr(t, "name", ""))
            for t in tools
        ]
        assert "local/robot-test" in names, f"Expected tool in {names}"
        print("tool-persist-add-then-list-ok")
    finally:
        os.unlink(db_path)


def fresh_list_empty() -> None:
    """List on a fresh (empty) database should return no tools."""
    fd, db_path = tempfile.mkstemp(suffix=".db", prefix="tool_persist_robot_")
    os.close(fd)
    try:
        eng: Engine = _engine(db_path)
        Base.metadata.create_all(eng)
        repo: ToolRegistryRepository = ToolRegistryRepository(
            session_factory=_factory(eng),
        )
        tools: list[Any] = repo.list_all()
        eng.dispose()

        assert len(tools) == 0, f"Expected 0 tools, got {len(tools)}"
        print("tool-persist-fresh-list-empty-ok")
    finally:
        os.unlink(db_path)


_COMMANDS: dict[str, Any] = {
    "add-then-list": add_then_list,
    "fresh-list-empty": fresh_list_empty,
}

if __name__ == "__main__":
    cmd: str = sys.argv[1] if len(sys.argv) > 1 else ""
    fn = _COMMANDS.get(cmd)
    if fn is None:
        print(
            f"Unknown command: {cmd!r}. Available: {list(_COMMANDS)}", file=sys.stderr
        )
        sys.exit(1)
    fn()