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..3b4c8cda5c 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() @@ -378,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]) +@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 + *, file_memory: SQLiteMemory, method: str, expire_after_busy: bool, started: float ) -> None: original = AsyncSession.execute calls: dict[Executable, int] = {} sessions: set[AsyncSession] = set() + 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 == pytest.approx(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