Spaces:
Running
Running
File size: 4,268 Bytes
6918275 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | import asyncio
import threading
import uuid
import weakref
import pytest
from helpers.defer import DeferredTask
class Owner:
pass
def make_task() -> DeferredTask:
return DeferredTask(f"defer-lifecycle-{uuid.uuid4()}")
def test_completed_task_releases_call_references_and_children():
task = make_task()
owner = Owner()
owner_ref = weakref.ref(owner)
child_killed = threading.Event()
class Child:
def kill(self, terminate_thread: bool = False) -> None:
assert terminate_thread
child_killed.set()
async def run(captured_owner):
return "done"
try:
task.add_child_task(Child(), terminate_thread=True) # type: ignore[arg-type]
task.start_task(run, owner)
assert task.result_sync(timeout=2) == "done"
assert child_killed.wait(2)
assert task.func is None
assert task.args == ()
assert task.kwargs == {}
del owner
assert owner_ref() is None
assert task.result_sync(timeout=2) == "done"
with pytest.raises(RuntimeError, match="Completed task cannot be restarted"):
task.restart()
finally:
task.kill(terminate_thread=True)
def test_run_task_end_extension_marks_state_dirty_after_completion(monkeypatch):
from extensions.python._functions.agent.AgentContext.run_task.end import (
_10_mark_state_dirty as task_done_extension,
)
task = make_task()
callback_called = threading.Event()
observations: list[tuple[str | None, bool]] = []
def mark_dirty(*, reason=None):
observations.append((reason, bool(task.is_alive())))
callback_called.set()
monkeypatch.setattr(
task_done_extension,
"mark_dirty_all",
mark_dirty,
)
async def run():
return "done"
try:
with pytest.raises(RuntimeError, match="Task hasn't been started"):
task.add_done_callback(lambda _future: None)
task.start_task(run)
task_done_extension.MarkStateDirty(agent=None).execute(
data={"result": task}
)
assert task.result_sync(timeout=2) == "done"
assert callback_called.wait(2)
assert observations == [("agent.AgentContext.run_task_done", False)]
finally:
task.kill(terminate_thread=True)
def test_kill_clears_stored_call_without_clearing_running_arguments():
task = make_task()
owner = Owner()
owner_ref = weakref.ref(owner)
started = threading.Event()
cancelled = threading.Event()
finished = threading.Event()
release: list[asyncio.Event] = []
async def run(captured_owner):
release.append(asyncio.Event())
started.set()
try:
await asyncio.Future()
except asyncio.CancelledError:
cancelled.set()
await release[0].wait()
finally:
finished.set()
try:
task.start_task(run, owner)
assert started.wait(2)
task.kill()
assert cancelled.wait(2)
assert task.func is None
assert task.args == ()
assert task.kwargs == {}
del owner
assert owner_ref() is not None
task.event_loop_thread.loop.call_soon_threadsafe(release[0].set)
assert finished.wait(2)
asyncio.run_coroutine_threadsafe(
asyncio.sleep(0), task.event_loop_thread.loop
).result(2)
assert owner_ref() is None
finally:
if release and task.event_loop_thread.loop:
task.event_loop_thread.loop.call_soon_threadsafe(release[0].set)
task.kill(terminate_thread=True)
def test_active_task_can_restart_from_its_snapshot():
task = make_task()
starts = [threading.Event(), threading.Event()]
run_count = 0
async def run(value):
nonlocal run_count
current_run = run_count
run_count += 1
assert value == "argument"
starts[current_run].set()
await asyncio.Future()
try:
task.start_task(run, "argument")
assert starts[0].wait(2)
task.restart()
assert starts[1].wait(2)
assert task.func is run
assert task.args == ("argument",)
finally:
task.kill(terminate_thread=True)
|