45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
|
"""Tests for the scan-wide budget-stop signal on the agent coordinator."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from strix.core.agents import AgentCoordinator
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_budget_stop_sets_flag() -> None:
|
||
|
|
coordinator = AgentCoordinator()
|
||
|
|
await coordinator.register("root", "strix", parent_id=None)
|
||
|
|
|
||
|
|
assert coordinator.budget_stopped is False
|
||
|
|
await coordinator.trigger_budget_stop()
|
||
|
|
assert coordinator.budget_stopped is True
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_budget_stop_unblocks_parked_agent() -> None:
|
||
|
|
# A parent parked in wait_for_message (awaiting a child) must be released so
|
||
|
|
# it can exit, no matter where in the tree the budget limit was hit.
|
||
|
|
coordinator = AgentCoordinator()
|
||
|
|
await coordinator.register("parent", "strix", parent_id=None)
|
||
|
|
|
||
|
|
waiter = asyncio.create_task(coordinator.wait_for_message("parent"))
|
||
|
|
await asyncio.sleep(0) # let the waiter park
|
||
|
|
assert not waiter.done()
|
||
|
|
|
||
|
|
await coordinator.trigger_budget_stop()
|
||
|
|
await asyncio.wait_for(waiter, timeout=1.0)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_wait_for_message_returns_immediately_after_budget_stop() -> None:
|
||
|
|
coordinator = AgentCoordinator()
|
||
|
|
await coordinator.register("agent", "recon", parent_id="parent")
|
||
|
|
await coordinator.trigger_budget_stop()
|
||
|
|
|
||
|
|
# No pending messages, but the stop flag short-circuits the wait.
|
||
|
|
await asyncio.wait_for(coordinator.wait_for_message("agent"), timeout=1.0)
|