mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-09-14 12:12:20 +02:00
229 lines
8.7 KiB
Python
229 lines
8.7 KiB
Python
import asyncio
|
|
import sqlite3
|
|
import threading
|
|
import pytest
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
from sqlalchemy import Column, DateTime, String, Text, create_engine
|
|
from sqlalchemy.orm import declarative_base, sessionmaker
|
|
|
|
|
|
def _setup_db(tmp_path, monkeypatch):
|
|
import core.database as cd
|
|
|
|
base = declarative_base()
|
|
|
|
class ScheduledTask(base):
|
|
__tablename__ = "scheduled_tasks"
|
|
|
|
id = Column(String, primary_key=True)
|
|
owner = Column(String)
|
|
name = Column(String)
|
|
task_type = Column(String, default="llm")
|
|
action = Column(String)
|
|
status = Column(String, default="active")
|
|
next_run = Column(DateTime)
|
|
last_run = Column(DateTime)
|
|
prompt = Column(Text, default='')
|
|
trigger_type = Column(String, default='schedule')
|
|
schedule = Column(String, default='daily')
|
|
scheduled_time = Column(String, default='08:00')
|
|
scheduled_day = Column(String)
|
|
scheduled_date = Column(DateTime)
|
|
cron_expression = Column(String)
|
|
|
|
class TaskRun(base):
|
|
__tablename__ = "task_runs"
|
|
|
|
id = Column(String, primary_key=True)
|
|
task_id = Column(String)
|
|
started_at = Column(DateTime)
|
|
finished_at = Column(DateTime)
|
|
status = Column(String)
|
|
result = Column(Text)
|
|
error = Column(Text)
|
|
model = Column(String)
|
|
|
|
engine = create_engine(f"sqlite:///{tmp_path / 'tasks.db'}")
|
|
base.metadata.create_all(engine)
|
|
session_local = sessionmaker(bind=engine, autocommit=False, autoflush=False)
|
|
monkeypatch.setattr(cd, "SessionLocal", session_local)
|
|
monkeypatch.setattr(cd, "ScheduledTask", ScheduledTask)
|
|
monkeypatch.setattr(cd, "TaskRun", TaskRun)
|
|
return session_local, ScheduledTask, TaskRun
|
|
|
|
|
|
def test_stop_task_cleans_up_queued_handle_and_run(tmp_path, monkeypatch):
|
|
session_local, ScheduledTask, TaskRun = _setup_db(tmp_path, monkeypatch)
|
|
|
|
db = session_local()
|
|
db.add(ScheduledTask(
|
|
id="queued-task",
|
|
owner="alice",
|
|
name="Queued Task",
|
|
task_type="llm",
|
|
status="active",
|
|
))
|
|
db.commit()
|
|
db.close()
|
|
|
|
from src.task_scheduler import TaskScheduler
|
|
|
|
async def drive():
|
|
scheduler = TaskScheduler.__new__(TaskScheduler)
|
|
scheduler._executing = {"queued-task"}
|
|
scheduler._executing_lock = asyncio.Lock()
|
|
scheduler._run_semaphore = asyncio.Semaphore(1)
|
|
scheduler._task_handles = {}
|
|
scheduler._concurrency_cap = 1
|
|
scheduler._task_defer_counts = {}
|
|
await scheduler._run_semaphore.acquire()
|
|
|
|
task = asyncio.create_task(scheduler._execute_task("queued-task"))
|
|
try:
|
|
for _ in range(50):
|
|
if "queued-task" in scheduler._task_handles:
|
|
db2 = session_local()
|
|
try:
|
|
run = db2.query(TaskRun).filter(TaskRun.task_id == "queued-task").first()
|
|
if run:
|
|
break
|
|
finally:
|
|
db2.close()
|
|
await asyncio.sleep(0.01)
|
|
else:
|
|
raise AssertionError("queued run was not created")
|
|
|
|
assert await scheduler.stop_task("queued-task") is True
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
finally:
|
|
scheduler._run_semaphore.release()
|
|
|
|
assert "queued-task" not in scheduler._task_handles
|
|
assert "queued-task" not in scheduler._executing
|
|
|
|
asyncio.run(drive())
|
|
|
|
db = session_local()
|
|
try:
|
|
run = db.query(TaskRun).filter(TaskRun.task_id == "queued-task").first()
|
|
assert run.status == "aborted"
|
|
assert run.error == "Stopped by user"
|
|
assert run.finished_at is not None
|
|
assert run.finished_at >= run.started_at
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
@pytest.mark.parametrize('mode', ['foreground', 'single', 'queued', 'queued_repeat'])
|
|
@pytest.mark.parametrize('lock_mode', ['EXCLUSIVE', 'IMMEDIATE'])
|
|
def test_cancel_keeps_event_loop_responsive_during_database_lock(tmp_path, monkeypatch, mode, lock_mode):
|
|
session_local, ScheduledTask, TaskRun = _setup_db(tmp_path, monkeypatch)
|
|
with session_local() as db:
|
|
if mode.startswith('queued'):
|
|
db.add(ScheduledTask(id='locked-task', owner='alice', name='Fixture',
|
|
task_type='llm', status='active',
|
|
next_run=datetime.now(timezone.utc).replace(tzinfo=None) - timedelta(minutes=1)))
|
|
else:
|
|
db.add(TaskRun(id='locked-run', task_id='locked-task', status='running'))
|
|
db.commit()
|
|
from src.task_scheduler import TaskScheduler
|
|
|
|
holder = sqlite3.connect(tmp_path / 'tasks.db', check_same_thread=False)
|
|
assert holder.execute('PRAGMA journal_mode').fetchone()[0] == 'wal'
|
|
heartbeat = threading.Event()
|
|
progress_before_release = []
|
|
def release_lock():
|
|
progress_before_release.append(heartbeat.is_set())
|
|
holder.commit()
|
|
release = threading.Timer(0.25, release_lock)
|
|
|
|
async def drive():
|
|
scheduler = TaskScheduler(None)
|
|
scheduler._executing.add('locked-task')
|
|
pending = None
|
|
if mode.startswith('queued'):
|
|
await scheduler._run_semaphore.acquire()
|
|
pending = asyncio.create_task(scheduler._execute_task('locked-task'))
|
|
await asyncio.sleep(0)
|
|
assert 'locked-task' in scheduler._task_handles
|
|
holder.execute(f'BEGIN {lock_mode}')
|
|
asyncio.get_running_loop().call_later(0.02, heartbeat.set)
|
|
repeated_stops = []
|
|
if mode == 'queued_repeat':
|
|
asyncio.get_running_loop().call_later(0.04, lambda: repeated_stops.append(
|
|
asyncio.create_task(scheduler.stop_task('locked-task'))))
|
|
release.start()
|
|
if mode == 'foreground':
|
|
assert await scheduler.stop_background_tasks_for_foreground() == 1
|
|
else:
|
|
assert await scheduler.stop_task('locked-task') is True
|
|
if pending:
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await pending
|
|
scheduler._run_semaphore.release()
|
|
if repeated_stops:
|
|
await asyncio.gather(*repeated_stops)
|
|
|
|
try:
|
|
asyncio.run(drive())
|
|
finally:
|
|
release.join(timeout=2)
|
|
holder.close()
|
|
assert progress_before_release == [True], 'foreground event loop froze behind SQLite'
|
|
with session_local() as db:
|
|
assert db.query(TaskRun).filter_by(task_id='locked-task').one().status == 'aborted'
|
|
if mode.startswith('queued'):
|
|
task = db.get(ScheduledTask, 'locked-task')
|
|
assert task.status == 'active'
|
|
assert task.next_run > datetime.now(timezone.utc).replace(tzinfo=None)
|
|
|
|
|
|
def test_running_task_cancel_keeps_event_loop_responsive_during_database_lock(tmp_path, monkeypatch):
|
|
session_local, ScheduledTask, TaskRun = _setup_db(tmp_path, monkeypatch)
|
|
from src.task_scheduler import TaskScheduler
|
|
from src.builtin_actions import BUILTIN_ACTIONS
|
|
monkeypatch.setenv('BACKGROUND_TASK_FOREGROUND_GATE', 'false')
|
|
with session_local() as db:
|
|
db.add(ScheduledTask(id='running-task', owner='alice', name='Fixture action',
|
|
task_type='action', action='fixture_wait', status='active'))
|
|
db.commit()
|
|
started = asyncio.Event()
|
|
async def external_action(**kwargs):
|
|
started.set()
|
|
await asyncio.Event().wait()
|
|
monkeypatch.setitem(BUILTIN_ACTIONS, 'fixture_wait', external_action)
|
|
holder = sqlite3.connect(tmp_path / 'tasks.db', check_same_thread=False)
|
|
heartbeat = threading.Event()
|
|
progress_before_release = []
|
|
def release_lock():
|
|
progress_before_release.append(heartbeat.is_set())
|
|
holder.commit()
|
|
release = threading.Timer(0.25, release_lock)
|
|
|
|
async def drive():
|
|
scheduler = TaskScheduler(None)
|
|
assert await scheduler.run_task_now('running-task')
|
|
await asyncio.wait_for(started.wait(), timeout=2)
|
|
pending = scheduler._task_handles['running-task']
|
|
holder.execute('BEGIN EXCLUSIVE')
|
|
asyncio.get_running_loop().call_later(0.02, heartbeat.set)
|
|
release.start()
|
|
assert await scheduler.stop_task('running-task')
|
|
await pending
|
|
try:
|
|
asyncio.run(drive())
|
|
finally:
|
|
if release.ident is not None:
|
|
release.join(timeout=2)
|
|
holder.close()
|
|
assert progress_before_release == [True], 'running-task cleanup blocked foreground'
|
|
with session_local() as db:
|
|
assert db.query(TaskRun).filter_by(task_id='running-task').one().status == 'aborted'
|
|
task = db.get(ScheduledTask, 'running-task')
|
|
assert task.status == 'active'
|
|
assert task.next_run > datetime.now(timezone.utc).replace(tzinfo=None)
|