EvoScientist-Multi/tests/test_worker_exit_runtime_compat.py

import asyncio
from concurrent.futures import Future

import pytest

from EvoScientist.langgraph_dev import worker_exit


@pytest.mark.asyncio
async def test_remote_thread_future_is_drained_before_cancellation_returns():
    remote = Future()
    entered = asyncio.Event()

    async def wait_remote():
        entered.set()
        return await worker_exit._await_remote_future(remote)

    task = asyncio.create_task(wait_remote())
    await entered.wait()
    task.cancel()
    await asyncio.sleep(0)
    assert not task.done()
    remote.set_result("finished")
    with pytest.raises(asyncio.CancelledError):
        await task


@pytest.mark.asyncio
async def test_cancellation_listener_restarts_after_idle_timeout(monkeypatch):
    calls = 0
    done = asyncio.Event()

    async def original(queue, run_id, thread_id, event):
        nonlocal calls
        calls += 1
        if calls == 2:
            event.set()

    await worker_exit._persistent_cancellation_listener(
        original, asyncio.Queue(), "run", "thread", done
    )
    assert calls == 2


@pytest.mark.asyncio
async def test_wait_for_exit_does_not_claim_a_still_active_worker(monkeypatch):
    monkeypatch.setattr(worker_exit, "cancel_and_inspect", lambda *_: {
        "execution_exited": False,
    })
    receipt = await worker_exit.wait_for_exit("thread", "run", timeout=0.01)
    assert receipt["execution_exited"] is False