From baaecfc28d941e81a97d2f849097575ee16011fa Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Sat, 10 Oct 2026 23:40:59 -0700 Subject: [PATCH 1/3] TEST: Make SQLite analytics lock checks scheduler-independent Drive lock deadline expiry from observed SQLITE_BUSY and a controlled request clock. Reuse isolated file memory, verify the exact reader connection's settings, arm cancellation budgets after writer setup, and observe completion without injecting cancellation. Preserve production retry semantics and semantic-report budgets. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- tests/unit/memory/test_attack_analytics.py | 42 ------- .../test_attack_analytics_lock_retry.py | 104 +++++++++++++++--- 2 files changed, 87 insertions(+), 59 deletions(-) diff --git a/tests/unit/memory/test_attack_analytics.py b/tests/unit/memory/test_attack_analytics.py index 6e053c4334..d8ff984804 100644 --- a/tests/unit/memory/test_attack_analytics.py +++ b/tests/unit/memory/test_attack_analytics.py @@ -540,48 +540,6 @@ def stop_during_query() -> int: assert (await reader.report_async(query=AttackAnalyticsQuery(), control=control())).counts == {} -@pytest.mark.parametrize("method", ["report", "results", "facets"]) -async def test_sqlite_deadline_bounds_file_database_lock_wait_and_restores_timeout( - *, tmp_path: Path, method: str -) -> None: - memory = SQLiteMemory.__new__(SQLiteMemory) - with patch.object(memory, "cleanup"): - memory.__init__(db_path=tmp_path / "analytics-lock.sqlite", skip_schema_migration=True) - try: - engine = memory._get_async_engine() - async with engine.begin() as connection: - await connection.run_sync(Base.metadata.create_all) - async with engine.connect() as writer: - default_timeout = (await writer.exec_driver_sql("PRAGMA busy_timeout")).scalar_one() - await writer.exec_driver_sql("BEGIN EXCLUSIVE") - try: - started = time.monotonic() - reader = AttackAnalyticsReader(memory=memory) - budget = QueryControl(deadline=started + 0.08) - with pytest.raises(AnalyticsTimeoutException): - if method == "report": - await reader.report_async(query=AttackAnalyticsQuery(), control=budget) - elif method == "results": - await reader.results_async(query=AttackAnalyticsResultsQuery(), control=budget) - else: - await reader.facets_async( - query=AttackAnalyticsFacetQuery(dimension=AttackAnalyticsDimension(name="operation")), - control=budget, - ) - assert time.monotonic() - started < 1 - async with engine.connect() as reader_connection: - assert ( - await reader_connection.exec_driver_sql("PRAGMA busy_timeout") - ).scalar_one() == default_timeout - finally: - await writer.rollback() - assert ( - await AttackAnalyticsReader(memory=memory).report_async(query=AttackAnalyticsQuery(), control=control()) - ).counts == {} - finally: - await memory.dispose_engine_async() - - async def test_odbc_timeout_is_applied_before_each_cursor_is_created() -> None: class RecordingConnection(sqlite3.Connection): timeout = 0 diff --git a/tests/unit/memory/test_attack_analytics_lock_retry.py b/tests/unit/memory/test_attack_analytics_lock_retry.py index 9cc79a42b9..989134f4fb 100644 --- a/tests/unit/memory/test_attack_analytics_lock_retry.py +++ b/tests/unit/memory/test_attack_analytics_lock_retry.py @@ -16,7 +16,7 @@ from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession from pyrit.exceptions.analytics_exception import AnalyticsTimeoutException -from pyrit.memory import SQLiteMemory +from pyrit.memory import SQLiteMemory, query_control from pyrit.memory.attack_analytics import AttackAnalyticsReader from pyrit.memory.memory_models import Base from pyrit.memory.query_control import QueryControl @@ -28,6 +28,7 @@ AttackOutcome, AttackResult, ) +from unit.async_utils import wait_for_completion_async if TYPE_CHECKING: from collections.abc import AsyncGenerator, Awaitable, Callable @@ -183,7 +184,9 @@ async def test_non_sqlite_execution_does_not_retry_sqlite_error_codes_async() -> retry.assert_not_awaited() -def _observe_busy(*, reader: AttackAnalyticsReader, observed: asyncio.Event) -> Callable[..., Awaitable[_ResultT]]: +def _observe_busy( + *, reader: AttackAnalyticsReader, observed: asyncio.Event, release: asyncio.Event | None = None +) -> Callable[..., Awaitable[_ResultT]]: retry = reader._retry_sqlite_busy_async async def observe_async(*, operation: Callable[[], Awaitable[_ResultT]], control: QueryControl) -> _ResultT: @@ -193,6 +196,8 @@ async def observed_operation_async() -> _ResultT: except OperationalError as error: if getattr(error.orig, "sqlite_errorcode", None) == sqlite3.SQLITE_BUSY: observed.set() + if release is not None: + await release.wait() raise return await retry(operation=observed_operation_async, control=control) @@ -200,6 +205,67 @@ async def observed_operation_async() -> _ResultT: return observe_async +@pytest.mark.parametrize("method", ["report", "matrix", "compact_report", "results", "facets"]) +async def test_sqlite_deadline_bounds_file_database_lock_wait_and_restores_timeout_async( + *, file_memory: SQLiteMemory, method: str +) -> None: + reader = AttackAnalyticsReader(memory=file_memory) + observed, release = asyncio.Event(), asyncio.Event() + engine = file_memory._get_async_engine() + async with engine.connect() as writer: + default_timeout = (await writer.exec_driver_sql("PRAGMA busy_timeout")).scalar_one() + writer_driver = (await writer.get_raw_connection()).driver_connection + await writer.exec_driver_sql("BEGIN EXCLUSIVE") + started = time.monotonic() + control = QueryControl(deadline=started + 10) + request_driver: SQLiteConnection | None = None + configure_timeout = reader._sqlite_busy_timeout_async + + async def configure_timeout_async(*, driver: SQLiteConnection, timeout_ms: int) -> None: + nonlocal request_driver + await configure_timeout(driver=driver, timeout_ms=timeout_ms) + if timeout_ms == 0: + request_driver = driver + # Arm the lock budget after acquisition, independently of test scheduling. + control.deadline = started + 0.08 + + with ( + patch.object(query_control, "time") as clock, + patch.object(reader, "_sqlite_busy_timeout_async", side_effect=configure_timeout_async), + patch.object( + reader, + "_retry_sqlite_busy_async", + side_effect=_observe_busy(reader=reader, observed=observed, release=release), + ), + ): + clock.monotonic.return_value = started + task = asyncio.create_task(_read_async(reader=reader, method=method, control=control)) + try: + await asyncio.wait_for(observed.wait(), timeout=30) + assert not task.done() + assert request_driver is not None and request_driver is not writer_driver + async with request_driver.execute("PRAGMA busy_timeout") as cursor: + assert await cursor.fetchone() == (0,) + assert control.remaining == pytest.approx(0.08) + clock.monotonic.return_value = control.deadline + assert control.expired and not control.cancel_event.is_set() + release.set() + with pytest.raises(AnalyticsTimeoutException): + await wait_for_completion_async(future=task) + assert not task.cancelled() and task.cancelling() == 0 + async with engine.connect() as connection: + assert (await connection.get_raw_connection()).driver_connection is request_driver + assert (await connection.exec_driver_sql("PRAGMA busy_timeout")).scalar_one() == default_timeout + assert (await writer.exec_driver_sql("PRAGMA busy_timeout")).scalar_one() == default_timeout + finally: + release.set() + await writer.rollback() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await _read_async(reader=reader, method=method, control=QueryControl(deadline=time.monotonic() + 10)) + + @pytest.mark.parametrize("method", ["report", "matrix", "compact_report", "results", "facets"]) async def test_sqlite_read_succeeds_when_writer_releases_lock_async(*, file_memory: SQLiteMemory, method: str) -> None: reader = AttackAnalyticsReader(memory=file_memory) @@ -215,16 +281,17 @@ async def test_sqlite_read_succeeds_when_writer_releases_lock_async(*, file_memo _read_async(reader=reader, method=method, control=QueryControl(deadline=time.monotonic() + 10)) ) try: - await asyncio.wait_for(observed.wait(), timeout=5) + await asyncio.wait_for(observed.wait(), timeout=30) assert not task.done() await writer.rollback() - await asyncio.wait_for(task, timeout=5) + await wait_for_completion_async(future=task) async with engine.connect() as connection: assert (await connection.exec_driver_sql("PRAGMA busy_timeout")).scalar_one() == default_timeout finally: - task.cancel() - await asyncio.gather(task, return_exceptions=True) await writer.rollback() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) @pytest.mark.parametrize("method", ["report", "matrix", "compact_report", "results", "facets"]) @@ -234,31 +301,34 @@ async def test_sqlite_lock_retry_cancellation_restores_connection_async( ) -> None: reader = AttackAnalyticsReader(memory=file_memory) observed = asyncio.Event() - control = QueryControl(deadline=time.monotonic() + 10) engine = file_memory._get_async_engine() async with engine.connect() as writer: default_timeout = (await writer.exec_driver_sql("PRAGMA busy_timeout")).scalar_one() await writer.exec_driver_sql("BEGIN EXCLUSIVE") + control = QueryControl(deadline=time.monotonic() + 10) with patch.object( reader, "_retry_sqlite_busy_async", side_effect=_observe_busy(reader=reader, observed=observed) ): task = asyncio.create_task(_read_async(reader=reader, method=method, control=control)) try: - await asyncio.wait_for(observed.wait(), timeout=5) + await asyncio.wait_for(observed.wait(), timeout=30) + assert not task.done() if cancel == "control": control.cancel() with pytest.raises(AnalyticsTimeoutException): - await asyncio.wait_for(task, timeout=1) + await wait_for_completion_async(future=task) + assert not task.cancelled() and task.cancelling() == 0 else: - task.cancel() - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(task, timeout=1) + task.cancel("requested cancellation") + with pytest.raises(asyncio.CancelledError, match="requested cancellation"): + await wait_for_completion_async(future=task) async with engine.connect() as connection: assert (await connection.exec_driver_sql("PRAGMA busy_timeout")).scalar_one() == default_timeout finally: - task.cancel() - await asyncio.gather(task, return_exceptions=True) await writer.rollback() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) await _read_async(reader=reader, method=method, control=QueryControl(deadline=time.monotonic() + 10)) @@ -326,12 +396,12 @@ async def read_async() -> None: with patch.object(reader, "_sqlite_busy_timeout_async", side_effect=restore_async): task = asyncio.create_task(read_async()) try: - await asyncio.wait_for(entered.wait(), timeout=5) + await asyncio.wait_for(entered.wait(), timeout=30) if cancel_in_body: task.cancel("original cancellation") else: finish_body.set() - await asyncio.wait_for(restoring.wait(), timeout=5) + await asyncio.wait_for(restoring.wait(), timeout=30) for index in range(2): task.cancel( "original cancellation" if not cancel_in_body and index == 0 else "cancel restoration again" @@ -340,7 +410,7 @@ async def read_async() -> None: assert not task.done() release_restore.set() with pytest.raises(asyncio.CancelledError, match="original cancellation"): - await asyncio.wait_for(task, timeout=5) + await wait_for_completion_async(future=task) finally: release_restore.set() task.cancel() From 753dc1c1e96821c347e04235d70296c79d80f961 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Sun, 11 Oct 2026 12:52:42 -0700 Subject: [PATCH 2/3] TEST: Control analytics projection retry deadlines Keep the query clock independent of CI scheduling while retaining a real completion watchdog. Exercise deadline expiry after SQLITE_BUSY for every projection without retrying or changing sessions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdb6b6f4-e9db-4f0b-99fc-61bda4a43474 --- .../test_attack_analytics_lock_retry.py | 32 +++++++++++++++---- 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/tests/unit/memory/test_attack_analytics_lock_retry.py b/tests/unit/memory/test_attack_analytics_lock_retry.py index 989134f4fb..627e6f243e 100644 --- a/tests/unit/memory/test_attack_analytics_lock_retry.py +++ b/tests/unit/memory/test_attack_analytics_lock_retry.py @@ -448,25 +448,43 @@ async def test_sqlite_reset_failure_discards_connection_and_preserves_cancellati @pytest.mark.parametrize("method", ["report", "matrix", "compact_report", "results", "facets"]) +@pytest.mark.parametrize("expire_after_busy", [False, True]) async def test_every_analytics_projection_retries_busy_in_the_same_session_async( - *, file_memory: SQLiteMemory, method: str + *, file_memory: SQLiteMemory, method: str, expire_after_busy: bool ) -> None: original = AsyncSession.execute calls: dict[Executable, int] = {} sessions: set[AsyncSession] = set() + started = time.monotonic() + control = QueryControl(deadline=started + 10) async def execute_async(self: AsyncSession, statement: Executable) -> Result[Any]: sessions.add(self) calls[statement] = calls.get(statement, 0) + 1 if calls[statement] == 1: + if expire_after_busy: + clock.monotonic.return_value = control.deadline raise _sqlite_error(sqlite3.SQLITE_BUSY) return await original(self, statement) - with patch.object(AsyncSession, "execute", execute_async): - await _read_async( - reader=AttackAnalyticsReader(memory=file_memory), - method=method, - control=QueryControl(deadline=time.monotonic() + 10), + with patch.object(query_control, "time") as clock, patch.object(AsyncSession, "execute", execute_async): + # Only the query clock is controlled; the event loop watchdog remains real. + clock.monotonic.return_value = started + task = asyncio.create_task( + _read_async(reader=AttackAnalyticsReader(memory=file_memory), method=method, control=control) ) - assert calls and all(count == 2 for count in calls.values()) + try: + if expire_after_busy: + with pytest.raises(AnalyticsTimeoutException): + await wait_for_completion_async(future=task) + assert control.expired and not control.cancel_event.is_set() + else: + await wait_for_completion_async(future=task) + assert control.remaining == 10 + assert not task.cancelled() and task.cancelling() == 0 + finally: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + assert calls and all(count == (1 if expire_after_busy else 2) for count in calls.values()) assert len(sessions) == 1 From 09d50dfae1fb519278cf56a72a9ef6b2600e96fe Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Sun, 11 Oct 2026 13:10:02 -0700 Subject: [PATCH 3/3] TEST: Compare analytics budgets with float tolerance Cover exact and fractional query clock baselines, including the CI baseline that reproduces floating-point cancellation. Preserve retry, session, expiry, and watchdog assertions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdb6b6f4-e9db-4f0b-99fc-61bda4a43474 --- tests/unit/memory/test_attack_analytics_lock_retry.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/memory/test_attack_analytics_lock_retry.py b/tests/unit/memory/test_attack_analytics_lock_retry.py index 627e6f243e..3b4c8cda5c 100644 --- a/tests/unit/memory/test_attack_analytics_lock_retry.py +++ b/tests/unit/memory/test_attack_analytics_lock_retry.py @@ -449,13 +449,13 @@ async def test_sqlite_reset_failure_discards_connection_and_preserves_cancellati @pytest.mark.parametrize("method", ["report", "matrix", "compact_report", "results", "facets"]) @pytest.mark.parametrize("expire_after_busy", [False, True]) +@pytest.mark.parametrize("started", [0.0, 248.670460394], ids=["exact-clock", "fractional-clock"]) async def test_every_analytics_projection_retries_busy_in_the_same_session_async( - *, file_memory: SQLiteMemory, method: str, expire_after_busy: bool + *, file_memory: SQLiteMemory, method: str, expire_after_busy: bool, started: float ) -> None: original = AsyncSession.execute calls: dict[Executable, int] = {} sessions: set[AsyncSession] = set() - started = time.monotonic() control = QueryControl(deadline=started + 10) async def execute_async(self: AsyncSession, statement: Executable) -> Result[Any]: @@ -480,7 +480,7 @@ async def execute_async(self: AsyncSession, statement: Executable) -> Result[Any assert control.expired and not control.cancel_event.is_set() else: await wait_for_completion_async(future=task) - assert control.remaining == 10 + assert control.remaining == pytest.approx(10) assert not task.cancelled() and task.cancelling() == 0 finally: if not task.done():