hermes-agent/tests/gateway/test_update_cron_drain.py
teknium1 ecc6725855 fix(gateway,cron): reconcile #60612 + #60631 onto one drain surface
Keep #60631's get_running_job_ids() snapshot + _active_cron_job_count()
(import-guarded for minimal test doubles) as the single read path, and
retarget #60612's drain tests at it. Drops the redundant
cron_jobs_in_flight() helper so there is one surface, not two.
2026-07-07 22:15:04 -07:00

86 lines
2.7 KiB
Python

"""Regression tests for #60432.
``/update`` (and other gateway shutdown paths) must drain in-flight cron jobs
before ``process_registry.kill_all()`` runs in final-cleanup. Cron work runs on
a thread-pool worker and is tracked in ``cron.scheduler._running_job_ids``, not
in ``GatewayRunner._running_agents`` — so a zero-agent drain must still wait
for cron to finish (or time out and take the interrupt/kill path).
"""
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from tests.gateway.restart_test_helpers import make_restart_runner
@pytest.mark.asyncio
async def test_drain_active_agents_waits_for_in_flight_cron_jobs():
runner, _adapter = make_restart_runner()
runner._running_agents = {}
cron_count = [1]
def _cron_in_flight():
return frozenset(f"job-{i}" for i in range(cron_count[0]))
async def finish_cron():
await asyncio.sleep(0.15)
cron_count[0] = 0
with patch("cron.scheduler.get_running_job_ids", side_effect=_cron_in_flight):
task = asyncio.create_task(finish_cron())
_snapshot, timed_out = await runner._drain_active_agents(1.0)
await task
assert timed_out is False
assert _snapshot == {}
@pytest.mark.asyncio
async def test_drain_active_agents_times_out_when_cron_still_running():
runner, _adapter = make_restart_runner()
runner._running_agents = {}
with patch("cron.scheduler.get_running_job_ids", return_value=frozenset({"job-1"})):
_snapshot, timed_out = await runner._drain_active_agents(0.05)
assert timed_out is True
assert _snapshot == {}
@pytest.mark.asyncio
async def test_gateway_stop_waits_for_cron_before_final_tool_kill():
"""Graceful cron completion must finish before final-cleanup kill_all."""
runner, adapter = make_restart_runner()
runner._restart_drain_timeout = 1.0
cron_count = [1]
call_order: list[str] = []
def _cron_in_flight():
return frozenset(f"job-{i}" for i in range(cron_count[0]))
def _fake_kill_all(task_id=None):
call_order.append("kill_all")
return 0
async def finish_cron():
await asyncio.sleep(0.12)
cron_count[0] = 0
with (
patch("cron.scheduler.get_running_job_ids", side_effect=_cron_in_flight),
patch("gateway.status.remove_pid_file"),
patch("gateway.status.write_runtime_status"),
patch("agent.auxiliary_client.shutdown_cached_clients"),
patch("tools.process_registry.process_registry") as registry_mock,
):
registry_mock.kill_all.side_effect = _fake_kill_all
adapter.disconnect = AsyncMock()
cron_task = asyncio.create_task(finish_cron())
await runner.stop()
await cron_task
assert call_order == ["kill_all"]