cleveragents-core/robot/helper_plan_lifecycle_persistence.py

"""Helper script for plan_lifecycle_persistence.robot smoke test."""

from __future__ import annotations

import sys
from typing import Any

sys.path.insert(0, "src")
sys.path.insert(0, ".")

from sqlalchemy import create_engine, event
from sqlalchemy.orm import Session, sessionmaker

from cleveragents.application.services.plan_lifecycle_service import (
    PlanLifecycleService,
)
from cleveragents.config.settings import Settings
from cleveragents.infrastructure.database.models import Base
from cleveragents.infrastructure.database.repositories import (
    ActionRepository,
    LifecyclePlanRepository,
)
from cleveragents.infrastructure.database.unit_of_work import UnitOfWork


def main() -> None:
    engine = create_engine(
        "sqlite:///:memory:",
        echo=False,
        future=True,
        connect_args={"check_same_thread": False},
    )

    @event.listens_for(engine, "connect")
    def _fk(dbapi_conn: Any, _rec: Any) -> None:
        cursor = dbapi_conn.cursor()
        cursor.execute("PRAGMA foreign_keys=ON")
        cursor.close()

    Base.metadata.create_all(engine)
    sf: sessionmaker[Session] = sessionmaker(
        bind=engine,
        expire_on_commit=False,
        autoflush=False,
        autocommit=False,
        class_=Session,
    )

    # Build a lightweight UoW that uses our test engine
    uow = UnitOfWork.__new__(UnitOfWork)
    uow.database_url = "sqlite:///:memory:"
    uow._engine = engine
    uow._session_factory = sf
    uow._database_initialized = True
    uow._prompt_for_migration = None

    settings = Settings()
    svc = PlanLifecycleService(settings=settings, unit_of_work=uow)

    # --- Test 1: create action persists ---
    svc.create_action(
        name="local/robot-persist",
        description="Robot test action",
        definition_of_done="All robot tests pass",
        strategy_actor="openai/gpt-4",
        execution_actor="openai/gpt-4",
    )
    session = sf()
    repo = ActionRepository(session_factory=lambda: session)
    db_action = repo.get_by_name("local/robot-persist")
    session.close()
    assert db_action is not None, "FAIL: action not found in DB"
    print("PASS: action persisted")

    # --- Test 2: create plan persists ---
    plan = svc.use_action(action_name="local/robot-persist")
    plan_id = plan.identity.plan_id
    session2 = sf()
    plan_repo = LifecyclePlanRepository(session_factory=lambda: session2)
    db_plan = plan_repo.get(plan_id)
    session2.close()
    assert db_plan is not None, "FAIL: plan not found in DB"
    assert db_plan.phase.value == "strategize", "FAIL: wrong phase"
    print("PASS: plan persisted")

    # --- Test 3: phase transition persists ---
    svc.start_strategize(plan_id)
    svc.complete_strategize(plan_id)
    svc.execute_plan(plan_id)
    session3 = sf()
    plan_repo3 = LifecyclePlanRepository(session_factory=lambda: session3)
    db_plan3 = plan_repo3.get(plan_id)
    session3.close()
    assert db_plan3 is not None, "FAIL: plan not found after transition"
    assert db_plan3.phase.value == "execute", "FAIL: wrong phase after execute"
    print("PASS: phase transition persisted")

    print("PASS: plan_lifecycle_persistence smoke test")


if __name__ == "__main__":
    main()