cleveragents-core/robot/helper_concurrency_locks.py

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

from __future__ import annotations

import sys
from datetime import UTC, datetime, timedelta
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.lock_service import LockService
from cleveragents.core.exceptions import LockConflictError, LockExpiredError
from cleveragents.infrastructure.database.models import Base, LockModel


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

    svc = LockService(session_factory=sf)

    # Test 1: basic acquire
    assert svc.acquire("owner-1", "plan", "p1") is True
    print("PASS: acquire")

    # Test 2: re-entrant
    assert svc.acquire("owner-1", "plan", "p1") is True
    print("PASS: re-entrant acquire")

    # Test 3: conflict
    try:
        svc.acquire("owner-2", "plan", "p1")
        print("FAIL: expected LockConflictError")
        sys.exit(1)
    except LockConflictError:
        print("PASS: conflict detection")

    # Test 4: release
    assert svc.release("owner-1", "plan", "p1") is True
    print("PASS: release")

    # Test 5: acquire after release
    assert svc.acquire("owner-2", "plan", "p1") is True
    print("PASS: acquire after release")

    # Test 6: renew
    assert svc.renew("owner-2", "plan", "p1") is True
    print("PASS: renew")

    # Test 7: release_all_for_owner
    svc.acquire("owner-3", "plan", "p10")
    svc.acquire("owner-3", "project", "proj-1")
    count = svc.release_all_for_owner("owner-3")
    assert count == 2, f"Expected 2, got {count}"
    print("PASS: release_all_for_owner")

    # Test 8: cleanup_expired
    now = datetime.now(tz=UTC)
    expired = (now - timedelta(seconds=60)).isoformat()
    acquired = (now - timedelta(seconds=120)).isoformat()
    session = sf()
    session.add(
        LockModel(
            owner_id="old",
            resource_type="plan",
            resource_id="stale-1",
            acquired_at=acquired,
            expires_at=expired,
        )
    )
    session.commit()
    session.close()

    cleaned = svc.cleanup_expired()
    assert cleaned >= 1, f"Expected >=1 cleaned, got {cleaned}"
    print("PASS: cleanup_expired")

    # Test 9: count_stale_locks returns 0 after cleanup
    assert svc.count_stale_locks() == 0
    print("PASS: count_stale_locks")

    # Test 10: is_locked
    svc.acquire("owner-4", "plan", "check-me")
    assert svc.is_locked("plan", "check-me") is True
    assert svc.is_locked("plan", "not-locked") is False
    print("PASS: is_locked")

    # Test 11: renew expired raises LockExpiredError
    session2 = sf()
    session2.add(
        LockModel(
            owner_id="owner-5",
            resource_type="plan",
            resource_id="expired-renew",
            acquired_at=acquired,
            expires_at=expired,
        )
    )
    session2.commit()
    session2.close()
    try:
        svc.renew("owner-5", "plan", "expired-renew")
        print("FAIL: expected LockExpiredError")
        sys.exit(1)
    except LockExpiredError:
        print("PASS: renew expired raises error")

    print("\nPASS: concurrency_locks smoke test")


if __name__ == "__main__":
    main()