cleveragents-core/robot/helper_async_execution.py

"""Helper script for Robot Framework async execution smoke tests.

Usage:
    python robot/helper_async_execution.py <test-name>

Supported test names:
    status-transitions   Test job status state machine
    worker-execute       Test worker pickup and execution
    payload-roundtrip    Test payload serialization round-trip
    health-report        Test worker health report
    stuck-detection      Test stuck job detection
    job-cleanup          Test completed job cleanup
"""

from __future__ import annotations

import sys
import traceback
from datetime import UTC, datetime, timedelta
from pathlib import Path

# Ensure src is on the path
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from ulid import ULID

from cleveragents.application.services.async_worker import (
    AsyncWorker,
    AsyncWorkerConfig,
    InMemoryJobStore,
)
from cleveragents.domain.models.core.async_job import (
    AsyncJob,
    AsyncJobStatus,
    InvalidJobTransitionError,
    deserialize_job_payload,
    serialize_job_payload,
)

PLAN_ID = "01HXYZ01HXYZ01HXYZ01HXYZ01"


def _job_id() -> str:
    return str(ULID())


def test_status_transitions() -> None:
    """Test job status transitions follow the state machine."""
    # queued -> running -> succeeded
    job = AsyncJob(job_id=_job_id(), plan_id=PLAN_ID, phase="execute")
    assert job.status == AsyncJobStatus.QUEUED, "FAIL: initial status should be queued"
    job.mark_running("w1")
    assert job.status == AsyncJobStatus.RUNNING, "FAIL: status should be running"
    job.mark_succeeded()
    assert job.status == AsyncJobStatus.SUCCEEDED, "FAIL: status should be succeeded"
    assert job.is_terminal, "FAIL: succeeded job should be terminal"
    print("  CHECK queued->running->succeeded: pass")

    # queued -> running -> failed (with error message)
    job2 = AsyncJob(job_id=_job_id(), plan_id=PLAN_ID, phase="execute")
    job2.mark_running("w2")
    job2.mark_failed(error="test error")
    assert job2.status == AsyncJobStatus.FAILED, "FAIL: status should be failed"
    assert job2.error_message == "test error", "FAIL: error_message should be set"
    print("  CHECK queued->running->failed (with error): pass")

    # queued -> cancelled
    job3 = AsyncJob(job_id=_job_id(), plan_id=PLAN_ID, phase="apply")
    job3.mark_cancelled()
    assert job3.status == AsyncJobStatus.CANCELLED, "FAIL: status should be cancelled"
    print("  CHECK queued->cancelled: pass")

    # Invalid: queued -> succeeded
    job4 = AsyncJob(job_id=_job_id(), plan_id=PLAN_ID, phase="execute")
    try:
        job4.mark_succeeded()
        raise AssertionError("Should have raised InvalidJobTransitionError")
    except InvalidJobTransitionError:
        pass
    print("  CHECK queued->succeeded raises error: pass")

    print("status-transitions-ok")


def test_worker_execute() -> None:
    """Test worker pickup and execution."""
    store = InMemoryJobStore()
    config = AsyncWorkerConfig(enabled=True, max_workers=2, poll_interval=0.1)
    worker = AsyncWorker(config, store)

    job = AsyncJob(job_id=_job_id(), plan_id=PLAN_ID, phase="execute")
    store.add(job)

    worker.pickup_and_execute(job)
    assert job.status == AsyncJobStatus.SUCCEEDED
    assert worker.health.jobs_processed == 1
    print("worker-execute-ok")


def test_payload_roundtrip() -> None:
    """Test payload serialization round-trip."""
    payload = serialize_job_payload(PLAN_ID, "execute", extra={"key": "val"})
    data = deserialize_job_payload(payload)
    assert data["plan_id"] == PLAN_ID
    assert data["phase"] == "execute"
    assert data["schema_version"] == 1
    assert data["extra"]["key"] == "val"
    print("payload-roundtrip-ok")


def test_health_report() -> None:
    """Test worker health report."""
    store = InMemoryJobStore()
    config = AsyncWorkerConfig(enabled=True, max_workers=4, poll_interval=0.5)
    worker = AsyncWorker(config, store)

    report = worker.get_health_report()
    assert "worker_id" in report
    assert "config" in report
    assert report["config"]["max_workers"] == 4
    print("health-report-ok")


def test_stuck_detection() -> None:
    """Test stuck job detection."""
    store = InMemoryJobStore()
    config = AsyncWorkerConfig(
        enabled=True, max_workers=2, poll_interval=0.1, job_timeout=1
    )
    worker = AsyncWorker(config, store)

    job = AsyncJob(job_id=_job_id(), plan_id=PLAN_ID, phase="execute")
    store.add(job)
    job.mark_running("worker-stuck")
    job.last_heartbeat = datetime.now(UTC) - timedelta(seconds=10)
    job.started_at = datetime.now(UTC) - timedelta(seconds=10)
    store.update(job)

    stuck = worker.detect_stuck_jobs()
    assert len(stuck) == 1
    assert stuck[0].status == AsyncJobStatus.FAILED
    print("stuck-detection-ok")


def test_job_cleanup() -> None:
    """Test completed job cleanup."""
    store = InMemoryJobStore()
    config = AsyncWorkerConfig(
        enabled=True, max_workers=2, poll_interval=0.1, job_ttl=1
    )
    worker = AsyncWorker(config, store)

    job = AsyncJob(job_id=_job_id(), plan_id=PLAN_ID, phase="execute")
    store.add(job)
    job.mark_running("worker-old")
    job.mark_succeeded()
    job.finished_at = datetime.now(UTC) - timedelta(seconds=10)
    store.update(job)

    cleaned = worker.cleanup_completed_jobs()
    assert cleaned == 1
    assert store.get(job.job_id) is None
    print("job-cleanup-ok")


TESTS = {
    "status-transitions": test_status_transitions,
    "worker-execute": test_worker_execute,
    "payload-roundtrip": test_payload_roundtrip,
    "health-report": test_health_report,
    "stuck-detection": test_stuck_detection,
    "job-cleanup": test_job_cleanup,
}


def main() -> None:
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <test-name>", file=sys.stderr)
        print(f"Available: {', '.join(TESTS)}", file=sys.stderr)
        sys.exit(1)

    test_name = sys.argv[1]
    test_fn = TESTS.get(test_name)
    if test_fn is None:
        print(f"Unknown test: {test_name}", file=sys.stderr)
        sys.exit(1)

    try:
        test_fn()
    except Exception as exc:
        # Print detailed assertion context for Robot test diagnostics
        print(f"FAIL: {test_name}: {exc}", file=sys.stderr)
        traceback.print_exc(file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()