ajarbot/tests/test_scheduler.py

#!/usr/bin/env python3
"""Test the TaskScheduler system."""

import sys
import traceback
from datetime import datetime

from agent import Agent
from scheduled_tasks import TaskScheduler


def test_schedule_calculation() -> bool:
    """Test schedule time calculations."""
    print("=" * 60)
    print("Testing Schedule Calculations")
    print("=" * 60)

    agent = Agent(
        provider="claude",
        workspace_dir="./memory_workspace",
        enable_heartbeat=False,
    )
    scheduler = TaskScheduler(
        agent, config_file="config/scheduled_tasks.yaml",
    )

    test_schedules = [
        "hourly",
        "daily 08:00",
        "daily 18:00",
        "weekly mon 09:00",
        "weekly fri 17:00",
    ]

    now = datetime.now()
    print(
        f"\nCurrent time: "
        f"{now.strftime('%Y-%m-%d %H:%M:%S %A')}\n"
    )

    for schedule in test_schedules:
        try:
            next_run = scheduler._calculate_next_run(schedule)
            time_until = next_run - now
            hours_until = time_until.total_seconds() / 3600

            formatted = next_run.strftime("%Y-%m-%d %H:%M %A")
            print(f"{schedule:20} -> {formatted}")
            print(
                f"{'':20}    (in {hours_until:.1f} hours)"
            )
        except Exception as e:
            print(f"{schedule:20} -> ERROR: {e}")

    return True


def test_task_loading() -> bool:
    """Test loading tasks from config."""
    print("\n" + "=" * 60)
    print("Testing Task Loading")
    print("=" * 60)

    agent = Agent(
        provider="claude",
        workspace_dir="./memory_workspace",
        enable_heartbeat=False,
    )
    scheduler = TaskScheduler(
        agent, config_file="config/scheduled_tasks.yaml",
    )

    tasks = scheduler.list_tasks()

    print(f"\nLoaded {len(tasks)} task(s):\n")

    for i, task in enumerate(tasks, 1):
        print(f"{i}. {task['name']}")
        print(f"   Schedule: {task['schedule']}")
        print(f"   Enabled: {task['enabled']}")
        print(f"   Next run: {task['next_run']}")
        if task["send_to"]:
            print(f"   Send to: {task['send_to']}")
        print()

    return len(tasks) > 0


def test_manual_execution() -> bool:
    """Test manual task execution."""
    print("=" * 60)
    print("Testing Manual Task Execution")
    print("=" * 60)

    agent = Agent(
        provider="claude",
        workspace_dir="./memory_workspace",
        enable_heartbeat=False,
    )
    scheduler = TaskScheduler(
        agent, config_file="config/scheduled_tasks.yaml",
    )

    if not scheduler.tasks:
        print(
            "\nNo tasks configured. "
            "Create tasks in config/scheduled_tasks.yaml"
        )
        return False

    test_task = next(
        (t for t in scheduler.tasks if t.enabled),
        scheduler.tasks[0],
    )

    print(f"\nManually executing task: {test_task.name}")
    print(f"Prompt: {test_task.prompt[:100]}...")
    print("\nExecuting...\n")

    result = scheduler.run_task_now(test_task.name)
    print(f"\nResult: {result}")

    return True


def test_scheduler_status() -> bool:
    """Test scheduler status reporting."""
    print("\n" + "=" * 60)
    print("Testing Scheduler Status")
    print("=" * 60)

    agent = Agent(
        provider="claude",
        workspace_dir="./memory_workspace",
        enable_heartbeat=False,
    )
    scheduler = TaskScheduler(
        agent, config_file="config/scheduled_tasks.yaml",
    )

    print(f"\nScheduler running: {scheduler.running}")
    print(f"Config file: {scheduler.config_file}")
    print(f"Tasks loaded: {len(scheduler.tasks)}")
    print(f"Adapters registered: {len(scheduler.adapters)}")

    enabled_count = sum(
        1 for t in scheduler.tasks if t.enabled
    )
    print(
        f"Enabled tasks: "
        f"{enabled_count}/{len(scheduler.tasks)}"
    )

    return True


def main() -> bool:
    """Run all tests."""
    print("\nTaskScheduler Test Suite\n")

    tests = [
        ("Schedule Calculation", test_schedule_calculation),
        ("Task Loading", test_task_loading),
        ("Scheduler Status", test_scheduler_status),
        # Commented out - uses API tokens:
        # ("Manual Execution", test_manual_execution),
    ]

    results = []

    for test_name, test_func in tests:
        try:
            result = test_func()
            results.append((test_name, result))
        except Exception as e:
            print(f"\nERROR in {test_name}: {e}")
            traceback.print_exc()
            results.append((test_name, False))

    print("\n" + "=" * 60)
    print("Test Summary")
    print("=" * 60)

    for test_name, passed in results:
        status = "+ PASS" if passed else "x FAIL"
        print(f"  {status}: {test_name}")

    passed_count = sum(1 for _, p in results if p)
    total_count = len(results)

    print(f"\n{passed_count}/{total_count} tests passed")

    if passed_count == total_count:
        print("\nAll tests passed! Scheduler is ready to use.")
        print("\nNext steps:")
        print("  1. Edit config/scheduled_tasks.yaml")
        print("  2. Set enabled: true for tasks you want")
        print("  3. Add your channel IDs")
        print(
            "  4. Run: python example_bot_with_scheduler.py"
        )
    else:
        print("\nSome tests failed. Check the output above.")

    return passed_count == total_count


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)