68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from app.gateway.admin_audit_retention import sweep_admin_audit_retention
|
|
from app.gateway.admin_audit_retention_control import (
|
|
is_paused,
|
|
read_status,
|
|
set_paused,
|
|
write_status,
|
|
)
|
|
|
|
|
|
class _MetricsStore:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
async def purge_before(self, _cutoff, *, batch_size: int) -> int:
|
|
self.calls += 1
|
|
return batch_size
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sweep_stops_after_committed_batch_and_reports_progress() -> None:
|
|
store = _MetricsStore()
|
|
stop_requested = False
|
|
progress: list[dict] = []
|
|
|
|
async def report(value: dict) -> None:
|
|
nonlocal stop_requested
|
|
progress.append(value)
|
|
if value["deleted"]["llm_metrics"] > 0:
|
|
stop_requested = True
|
|
|
|
deleted = await sweep_admin_audit_retention(
|
|
session_factory=None,
|
|
llm_metrics_store=store,
|
|
retention_days=90,
|
|
batch_size=3,
|
|
max_batches_per_store=4,
|
|
inter_batch_pause_seconds=0,
|
|
max_sweep_seconds=5,
|
|
should_stop=lambda: stop_requested,
|
|
on_progress=report,
|
|
)
|
|
|
|
assert store.calls == 1
|
|
assert deleted == {
|
|
"llm_metrics": 3,
|
|
"tool_metrics": 0,
|
|
"run_summaries": 0,
|
|
"leaderboard_snapshots": 0,
|
|
}
|
|
assert progress[-1]["status"] == "paused"
|
|
assert progress[-1]["deleted"]["llm_metrics"] == 3
|
|
|
|
|
|
def test_pause_marker_and_progress_survive_process_restart(tmp_path) -> None:
|
|
assert not is_paused(tmp_path)
|
|
set_paused(tmp_path, True)
|
|
write_status(tmp_path, {"status": "running", "deleted": {"llm_metrics": 12}})
|
|
|
|
assert is_paused(tmp_path)
|
|
assert read_status(tmp_path) == {"status": "running", "deleted": {"llm_metrics": 12}}
|
|
|
|
set_paused(tmp_path, False)
|
|
assert not is_paused(tmp_path)
|