mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-17 02:09:15 +00:00
d2b5f884e3
Queue comments received while a run is active, then submit one deduplicated follow-up after it finishes. Failed drains are requeued and watcher tasks stop cleanly with the channel manager.
259 lines
11 KiB
Python
259 lines
11 KiB
Python
"""Regression tests for Gateway lifespan shutdown.
|
|
|
|
These tests guard the invariant that lifespan shutdown is *bounded*: a
|
|
misbehaving channel whose ``stop()`` blocks forever must not keep the
|
|
uvicorn worker alive. A hung worker is the precondition for the
|
|
signal-reentrancy deadlock described in
|
|
``app.gateway.app._SHUTDOWN_HOOK_TIMEOUT_SECONDS``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from contextlib import asynccontextmanager
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from fastapi import FastAPI
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _noop_langgraph_runtime(_app, _startup_config):
|
|
yield
|
|
|
|
|
|
async def _run_lifespan_with_hanging_stop() -> float:
|
|
"""Drive the lifespan context with stop_channel_service hanging forever.
|
|
|
|
Returns the elapsed wall-clock seconds.
|
|
"""
|
|
from app.gateway.app import _SHUTDOWN_HOOK_TIMEOUT_SECONDS, lifespan
|
|
|
|
async def hang_forever() -> None:
|
|
await asyncio.sleep(3600)
|
|
|
|
app = FastAPI()
|
|
startup_config = MagicMock()
|
|
startup_config.log_level = "INFO"
|
|
# Keep this test focused on the channel-hang timing: skip the memory drain.
|
|
startup_config.memory.enabled = False
|
|
startup_config.memory.shutdown_flush_timeout_seconds = 5.0
|
|
fake_service = MagicMock()
|
|
fake_service.get_status = MagicMock(return_value={})
|
|
|
|
async def fake_start(_startup_config, **_kwargs):
|
|
return fake_service
|
|
|
|
close_oidc_service = AsyncMock()
|
|
|
|
with (
|
|
patch("app.gateway.app.get_app_config", return_value=startup_config),
|
|
patch("app.gateway.app.get_gateway_config", return_value=MagicMock(host="x", port=0)),
|
|
patch("app.gateway.app.langgraph_runtime", _noop_langgraph_runtime),
|
|
patch("app.gateway.app.auth.close_oidc_service", close_oidc_service),
|
|
patch("app.channels.service.start_channel_service", side_effect=fake_start),
|
|
patch("app.channels.service.stop_channel_service", side_effect=hang_forever),
|
|
patch("deerflow.agents.memory.get_memory_manager", return_value=MagicMock()),
|
|
):
|
|
loop = asyncio.get_event_loop()
|
|
start = loop.time()
|
|
async with lifespan(app):
|
|
pass
|
|
elapsed = loop.time() - start
|
|
|
|
close_oidc_service.assert_awaited_once()
|
|
assert _SHUTDOWN_HOOK_TIMEOUT_SECONDS < 30.0, "Timeout constant must stay modest"
|
|
return elapsed
|
|
|
|
|
|
def test_shutdown_is_bounded_when_channel_stop_hangs():
|
|
"""Lifespan exit must complete near the configured timeout, not hang."""
|
|
from app.gateway.app import _SHUTDOWN_HOOK_TIMEOUT_SECONDS
|
|
|
|
elapsed = asyncio.run(_run_lifespan_with_hanging_stop())
|
|
|
|
# Generous upper bound: timeout + 2s slack for scheduling overhead.
|
|
assert elapsed < _SHUTDOWN_HOOK_TIMEOUT_SECONDS + 2.0, f"Lifespan shutdown took {elapsed:.2f}s; expected <= {_SHUTDOWN_HOOK_TIMEOUT_SECONDS + 2.0:.1f}s"
|
|
# Lower bound: the wait_for should actually have waited.
|
|
assert elapsed >= _SHUTDOWN_HOOK_TIMEOUT_SECONDS - 0.5, f"Lifespan exited too quickly ({elapsed:.2f}s); wait_for may not have been invoked."
|
|
|
|
|
|
async def _run_lifespan_with_upload_staging_cleanup():
|
|
from app.gateway.app import lifespan
|
|
|
|
app = FastAPI()
|
|
startup_config = SimpleNamespace(log_level="INFO", memory=SimpleNamespace(token_counting="char", enabled=False, shutdown_flush_timeout_seconds=30.0))
|
|
fake_service = MagicMock()
|
|
fake_service.get_status = MagicMock(return_value={})
|
|
cleanup_upload_staging_files = MagicMock(return_value=2)
|
|
close_oidc_service = AsyncMock()
|
|
stop_channel_service = AsyncMock()
|
|
|
|
async def fake_start(_startup_config, **_kwargs):
|
|
return fake_service
|
|
|
|
with (
|
|
patch("app.gateway.app.get_app_config", return_value=startup_config),
|
|
patch("app.gateway.app.get_gateway_config", return_value=MagicMock(host="x", port=0)),
|
|
patch("app.gateway.app.langgraph_runtime", _noop_langgraph_runtime),
|
|
patch("app.gateway.app.cleanup_stale_upload_staging_files", cleanup_upload_staging_files),
|
|
patch("app.gateway.app.auth.close_oidc_service", close_oidc_service),
|
|
patch("app.channels.service.start_channel_service", side_effect=fake_start),
|
|
patch("app.channels.service.stop_channel_service", stop_channel_service),
|
|
):
|
|
async with lifespan(app):
|
|
pass
|
|
|
|
return cleanup_upload_staging_files, close_oidc_service, stop_channel_service
|
|
|
|
|
|
def test_lifespan_sweeps_upload_staging_files_on_startup():
|
|
cleanup_upload_staging_files, close_oidc_service, stop_channel_service = asyncio.run(_run_lifespan_with_upload_staging_cleanup())
|
|
|
|
cleanup_upload_staging_files.assert_called_once_with()
|
|
close_oidc_service.assert_awaited_once()
|
|
stop_channel_service.assert_awaited_once()
|
|
|
|
|
|
async def _run_lifespan_with_memory_flush(*, enabled: bool, flush_return: bool) -> MagicMock:
|
|
"""Drive lifespan with a spied memory manager.shutdown_flush.
|
|
|
|
Returns the manager mock so the caller can assert the shutdown flush was
|
|
reached (and with what timeout). The host calls ``shutdown_flush``
|
|
unconditionally when memory is enabled -- there is no host-level
|
|
``pending_count/is_processing`` gate, because the backend short-circuits on
|
|
an idle buffer and keeping the in-flight race inside the backend means the
|
|
host cannot "forget" it (review #6 on the original PR).
|
|
"""
|
|
from app.gateway.app import lifespan
|
|
|
|
app = FastAPI()
|
|
startup_config = SimpleNamespace(
|
|
log_level="INFO",
|
|
memory=SimpleNamespace(
|
|
token_counting="char",
|
|
enabled=enabled,
|
|
shutdown_flush_timeout_seconds=5.0,
|
|
),
|
|
)
|
|
fake_service = MagicMock()
|
|
fake_service.get_status = MagicMock(return_value={})
|
|
close_oidc_service = AsyncMock()
|
|
stop_channel_service = AsyncMock()
|
|
|
|
async def fake_start(_startup_config, **_kwargs):
|
|
return fake_service
|
|
|
|
manager = MagicMock()
|
|
manager.shutdown_flush.return_value = flush_return
|
|
|
|
with (
|
|
patch("app.gateway.app.get_app_config", return_value=startup_config),
|
|
patch("app.gateway.app.get_gateway_config", return_value=MagicMock(host="x", port=0)),
|
|
patch("app.gateway.app.langgraph_runtime", _noop_langgraph_runtime),
|
|
patch("app.gateway.app.auth.close_oidc_service", close_oidc_service),
|
|
patch("app.channels.service.start_channel_service", side_effect=fake_start),
|
|
patch("app.channels.service.stop_channel_service", stop_channel_service),
|
|
patch("deerflow.agents.memory.get_memory_manager", return_value=manager),
|
|
):
|
|
async with lifespan(app):
|
|
pass
|
|
|
|
return manager
|
|
|
|
|
|
def test_lifespan_drains_memory_on_shutdown_with_configured_timeout(caplog) -> None:
|
|
"""When memory is enabled, shutdown calls manager.shutdown_flush with the
|
|
configured timeout (asserts the timeout is forwarded, review #3) and logs
|
|
'completed' at INFO when the drain finishes."""
|
|
caplog.set_level(logging.INFO, logger="app.gateway.app")
|
|
manager = asyncio.run(_run_lifespan_with_memory_flush(enabled=True, flush_return=True))
|
|
manager.shutdown_flush.assert_called_once_with(5.0)
|
|
assert any(r.levelno == logging.INFO and "flush completed" in r.message for r in caplog.records)
|
|
|
|
|
|
def test_lifespan_warns_when_memory_flush_does_not_finish(caplog) -> None:
|
|
"""A False return (timeout/failure) is the path operators actually see when
|
|
K8s SIGKILLs the drain; the host must log a WARNING (not 'completed'), so
|
|
the loss risk is visible (review #3 False-branch coverage; review #2/#4
|
|
failed-flush semantics)."""
|
|
caplog.set_level(logging.WARNING, logger="app.gateway.app")
|
|
manager = asyncio.run(_run_lifespan_with_memory_flush(enabled=True, flush_return=False))
|
|
manager.shutdown_flush.assert_called_once_with(5.0)
|
|
assert any(r.levelno == logging.WARNING and "did not finish" in r.message for r in caplog.records)
|
|
assert not any("flush completed" in r.message for r in caplog.records)
|
|
|
|
|
|
def test_lifespan_skips_memory_flush_when_disabled() -> None:
|
|
"""memory.enabled=False skips the drain entirely."""
|
|
manager = asyncio.run(_run_lifespan_with_memory_flush(enabled=False, flush_return=True))
|
|
manager.shutdown_flush.assert_not_called()
|
|
|
|
|
|
# ── startup warm-up log accuracy ────────────────────────────────────────────
|
|
|
|
|
|
async def _run_lifespan_with_warm_return(warm_return: bool | None) -> MagicMock:
|
|
"""Drive lifespan with a spied ``manager.warm`` returning ``warm_return``.
|
|
|
|
The startup warm block reads the tri-state return: None = nothing to warm
|
|
(logs "skipping"), True = warmed, False = failed (logs WARNING). Returns the
|
|
manager mock so the caller can assert warm was reached.
|
|
"""
|
|
from app.gateway.app import lifespan
|
|
|
|
app = FastAPI()
|
|
startup_config = SimpleNamespace(
|
|
log_level="INFO",
|
|
memory=SimpleNamespace(
|
|
token_counting="char",
|
|
enabled=False,
|
|
shutdown_flush_timeout_seconds=5.0,
|
|
),
|
|
)
|
|
fake_service = MagicMock()
|
|
fake_service.get_status = MagicMock(return_value={})
|
|
close_oidc_service = AsyncMock()
|
|
stop_channel_service = AsyncMock()
|
|
|
|
async def fake_start(_startup_config, **_kwargs):
|
|
return fake_service
|
|
|
|
manager = MagicMock()
|
|
manager.warm.return_value = warm_return
|
|
|
|
with (
|
|
patch("app.gateway.app.get_app_config", return_value=startup_config),
|
|
patch("app.gateway.app.get_gateway_config", return_value=MagicMock(host="x", port=0)),
|
|
patch("app.gateway.app.langgraph_runtime", _noop_langgraph_runtime),
|
|
patch("app.gateway.app.auth.close_oidc_service", close_oidc_service),
|
|
patch("app.channels.service.start_channel_service", side_effect=fake_start),
|
|
patch("app.channels.service.stop_channel_service", stop_channel_service),
|
|
patch("deerflow.agents.memory.get_memory_manager", return_value=manager),
|
|
):
|
|
async with lifespan(app):
|
|
pass
|
|
|
|
return manager
|
|
|
|
|
|
def test_lifespan_logs_skipping_when_backend_has_nothing_to_warm(caplog) -> None:
|
|
"""A backend whose warm() returns None (base default -- nothing to warm,
|
|
e.g. noop) logs "skipping" at INFO, not the misleading "warmed successfully"
|
|
(a non-DeerMem backend never touched the tiktoken cache)."""
|
|
caplog.set_level(logging.INFO, logger="app.gateway.app")
|
|
manager = asyncio.run(_run_lifespan_with_warm_return(None))
|
|
manager.warm.assert_called_once_with()
|
|
assert any(r.levelno == logging.INFO and "nothing to warm" in r.message for r in caplog.records)
|
|
assert not any("warmed successfully" in r.message for r in caplog.records)
|
|
|
|
|
|
def test_lifespan_warns_when_warm_returns_false(caplog) -> None:
|
|
"""warm()=False means warming was attempted and failed; the host logs a
|
|
WARNING so the operator sees the character-based-fallback degradation."""
|
|
caplog.set_level(logging.WARNING, logger="app.gateway.app")
|
|
manager = asyncio.run(_run_lifespan_with_warm_return(False))
|
|
manager.warm.assert_called_once_with()
|
|
assert any(r.levelno == logging.WARNING and "warm-up failed" in r.message for r in caplog.records)
|