diff --git a/tests/gateway/test_update_cron_drain.py b/tests/gateway/test_update_cron_drain.py new file mode 100644 index 00000000000..ee34d72cd82 --- /dev/null +++ b/tests/gateway/test_update_cron_drain.py @@ -0,0 +1,86 @@ +"""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 cron_count[0] + + async def finish_cron(): + await asyncio.sleep(0.15) + cron_count[0] = 0 + + with patch("cron.scheduler.cron_jobs_in_flight", 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.cron_jobs_in_flight", return_value=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 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.cron_jobs_in_flight", 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"]