From b8c1dfceb0d2b6315ac6fd35ec3ad3c6f789fdca Mon Sep 17 00:00:00 2001 From: jbolor21 <86250273+jbolor21@users.noreply.github.com> Date: Fri, 23 Oct 2026 12:17:52 -0700 Subject: [PATCH] reducing scenario db reads --- .../backend/services/scenario_run_service.py | 40 +-- .../unit/backend/test_scenario_run_service.py | 256 ++++++++++++++---- 2 files changed, 225 insertions(+), 71 deletions(-) diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 8f3328e34b..0e2b121dd6 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -450,7 +450,7 @@ async def _schedule_prepared_run_async( self, *, request: RunScenarioRequest, prepared: _PreparedRun, resumed_from_cancelled: bool ) -> ScenarioRunSummary: """ - Validate persisted preparation and admit it to the FIFO scheduler. + Validate the persisted header before admission; build details only for the returned response. Returns: ScenarioRunSummary: The persisted active, queued, or terminal state. @@ -461,10 +461,10 @@ async def _schedule_prepared_run_async( raise ValueError("Scenario did not produce a scenario_result_id during initialization.") if request.scenario_result_id and scenario_result_id != request.scenario_result_id: raise ValueError("Scenario initialization changed the saved result ID; resume was not started.") - persisted = await self._memory.get_scenario_results_async(scenario_result_ids=[scenario_result_id]) - if not persisted: + persisted = await self._memory.get_scenario_result_header_async(scenario_result_id=scenario_result_id) + if persisted is None: raise RuntimeError(f"Scenario run {scenario_result_id} was not persisted during initialization.") - if not resumed_from_cancelled and persisted[0].scenario_run_state == ScenarioRunState.CANCELLED: + if not resumed_from_cancelled and persisted.scenario_run_state == ScenarioRunState.CANCELLED: response = await self._build_response_async( scenario_result_id=scenario_result_id, active_error=None, @@ -476,23 +476,13 @@ async def _schedule_prepared_run_async( f"Scenario run {scenario_result_id} was not found in the database after initialization." ) return response - if ( - await self._build_response_async( - scenario_result_id=scenario_result_id, - active_error=None, - queue_position=None, - active_scenario_result_id=self._active_scenario_result_id, - ) - is None - ): - raise RuntimeError(f"Scenario run {scenario_result_id} was not found in the database after initialization.") scheduled = _ActiveTask( scenario_result_id=scenario_result_id, scenario=scenario, adversarial_target=prepared.adversarial_target, - scenario_name=persisted[0].scenario_name, + scenario_name=persisted.scenario_name, scenario_registry_name=request.scenario_name, - created_at=persisted[0].creation_time, + created_at=persisted.creation_time, enqueued_at=datetime.now(UTC), ) await self._enqueue_run_async(scheduled=scheduled) @@ -839,11 +829,11 @@ async def cancel_run_async(self, *, scenario_result_id: str) -> ScenarioRunSumma Raises: ValueError: If the run is already in a terminal state or not active. """ - results = await self._memory.get_scenario_results_async(scenario_result_ids=[scenario_result_id]) - if not results: + header = await self._memory.get_scenario_result_header_async(scenario_result_id=scenario_result_id) + if header is None: return None - db_status = results[0].scenario_run_state + db_status = header.scenario_run_state if self._is_terminal_state(db_status): raise ValueError(f"Cannot cancel run in '{db_status}' state.") @@ -870,9 +860,9 @@ async def cancel_run_async(self, *, scenario_result_id: str) -> ScenarioRunSumma active.cancellation_error_type = "CancelledError" task = active.task else: - latest = await self._memory.get_scenario_results_async(scenario_result_ids=[scenario_result_id]) - if latest and self._is_terminal_state(latest[0].scenario_run_state): - raise ValueError(f"Cannot cancel run in '{latest[0].scenario_run_state}' state.") + latest = await self._memory.get_scenario_result_header_async(scenario_result_id=scenario_result_id) + if latest is not None and self._is_terminal_state(latest.scenario_run_state): + raise ValueError(f"Cannot cancel run in '{latest.scenario_run_state}' state.") await self._memory.try_update_scenario_run_state_async( scenario_result_id=scenario_result_id, expected_states={ @@ -1093,10 +1083,10 @@ async def _handoff_scheduler_async(self, *, scenario_result_id: str) -> None: return while self._queued_runs: next_run = self._queued_runs[0] - persisted = await self._memory.get_scenario_results_async( - scenario_result_ids=[next_run.scenario_result_id] + persisted = await self._memory.get_scenario_result_header_async( + scenario_result_id=next_run.scenario_result_id ) - if not persisted or persisted[0].scenario_run_state != ScenarioRunState.QUEUED: + if persisted is None or persisted.scenario_run_state != ScenarioRunState.QUEUED: self._queued_runs.popleft() self._queue_revision += 1 continue diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index f29c7a64e1..e98dbb2a55 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -135,6 +135,36 @@ async def test_missing_operation_rejected_before_initializers_async(sqlite_insta await service.close_async() +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("invalid_field", ["scenario_identifier", "scenario_run_state"]) +async def test_invalid_persisted_header_blocks_execution_async( + *, sqlite_instance: SQLiteMemory, invalid_field: str +) -> None: + stored = make_scenario_result(objective_target_identifier=get_mock_target_identifier(), attack_results={}) + entry = ScenarioResultEntry(entry=stored) + if invalid_field == "scenario_identifier": + entry.scenario_identifier = {"class_name": 123} + else: + entry.scenario_run_state = "invalid-state" + with sqlite_instance._get_session() as session: + session.add(entry) + session.commit() + scenario = MagicMock(spec=Scenario) + scenario._scenario_result_id = str(stored.id) + service = ScenarioRunService() + try: + with ( + patch.object(service, "_prepare_run_async", return_value=_svc_mod._PreparedRun(scenario=scenario)), + patch.object(service, "_enqueue_run_async", new_callable=AsyncMock) as enqueue, + pytest.raises(ValueError), + ): + await service.start_run_async(request=_make_request()) + enqueue.assert_not_awaited() + scenario.run_async.assert_not_awaited() + finally: + await service.close_async() + + @pytest.mark.usefixtures("patch_central_database") @pytest.mark.parametrize("explicit_id", [False, True]) async def test_fresh_run_keeps_deprecated_name_unless_an_id_is_given_async(sqlite_instance, explicit_id: bool) -> None: @@ -469,6 +499,7 @@ def mock_memory(): """Patch CentralMemory.get_memory_instance to return a mock.""" mock = MagicMock(spec=SQLiteMemory) mock.get_scenario_results_async = AsyncMock(return_value=[]) + mock.get_scenario_result_header_async = AsyncMock(return_value=None) mock.get_scenario_run_history_page_async = AsyncMock(return_value=([], {}, False)) mock.get_scenario_history_aggregates_async = AsyncMock(return_value={}) # Default: no error AttackResults linked to any scenario. Tests that exercise @@ -509,6 +540,7 @@ def mock_all_registries(mock_memory): # By default, return a matching DB result for async run reads. db_result = _make_db_scenario_result() mock_memory.get_scenario_results_async = AsyncMock(return_value=[db_result]) + mock_memory.get_scenario_result_header_async = AsyncMock(return_value=db_result) with ( patch(f"{_REGISTRY_PATCH_BASE}.ScenarioRegistry.get_registry_singleton", return_value=mock_sr), @@ -608,6 +640,9 @@ def update_state(*, scenario_result_id: str, scenario_run_state: ScenarioRunStat mock_all_registries["scenario_registry"].create_and_initialize_async.side_effect = initialize_async memory = mock_all_registries["memory"] memory.get_scenario_results_async = AsyncMock(side_effect=get_results) + memory.get_scenario_result_header_async = AsyncMock( + side_effect=lambda *, scenario_result_id: records.get(scenario_result_id) + ) memory.update_scenario_run_state_async = AsyncMock(side_effect=update_state) memory.try_update_scenario_run_state_async = AsyncMock(side_effect=update_state) memory.update_scenario_run_state_and_metadata_fields_async = AsyncMock(side_effect=update_state) @@ -961,6 +996,9 @@ def _update_state(*, scenario_result_id: str, scenario_run_state: ScenarioRunSta mock_sr.create_and_initialize_async = AsyncMock(side_effect=_create_scenario) mock_memory.get_scenario_results_async = AsyncMock(side_effect=_get_results) + mock_memory.get_scenario_result_header_async = AsyncMock( + side_effect=lambda *, scenario_result_id: records.get(scenario_result_id) + ) mock_memory.update_scenario_run_state_async = AsyncMock(side_effect=_update_state) mock_memory.try_update_scenario_run_state_async = AsyncMock(side_effect=_update_state) mock_memory.update_scenario_run_state_and_metadata_fields_async = AsyncMock(side_effect=_update_state) @@ -1287,14 +1325,16 @@ async def _run_async() -> None: async def _get_results_async(*, scenario_result_ids: list[str] | None = None) -> list[MagicMock]: if scenario_result_ids is None: return list(records.values()) - if fail_once["handoff_read"] and scenario_result_ids == ["run-2"]: + return [records[run_id] for run_id in scenario_result_ids if run_id in records] + + async def _get_header_async(*, scenario_result_id: str) -> MagicMock | None: + if fail_once["handoff_read"] and scenario_result_id == "run-2": fail_once["handoff_read"] = False raise RuntimeError("temporary storage failure") - if len(scenario_result_ids) == 1 and scenario_result_ids[0] in blocked_handoffs: - run_id = scenario_result_ids[0] - handoff_read_started[run_id].set() - await release_handoff[run_id].wait() - return [records[run_id] for run_id in scenario_result_ids if run_id in records] + if scenario_result_id in blocked_handoffs: + handoff_read_started[scenario_result_id].set() + await release_handoff[scenario_result_id].wait() + return records.get(scenario_result_id) def _update_state(*, scenario_result_id: str, scenario_run_state: ScenarioRunState, **_: object) -> None: if ( @@ -1329,6 +1369,7 @@ async def _observe_handoff_async(run_id: str) -> None: mock_sr.create_and_initialize_async = AsyncMock(side_effect=_create_scenario_async) mock_memory.get_scenario_results_async = AsyncMock(side_effect=_get_results_async) + mock_memory.get_scenario_result_header_async = AsyncMock(side_effect=_get_header_async) mock_memory.update_scenario_run_state_async = AsyncMock(side_effect=_update_state) mock_memory.try_update_scenario_run_state_async = AsyncMock(side_effect=_update_state) mock_memory.update_scenario_run_state_and_metadata_fields_async = AsyncMock(side_effect=_update_state) @@ -1572,7 +1613,7 @@ async def _complete_then_cancel_async(awaitable): ScenarioRunState.QUEUED, } - @pytest.mark.parametrize("phase", ["persisted-read", "summary-read", "active-write", "queued-write"]) + @pytest.mark.parametrize("phase", ["persisted-read", "active-write", "queued-write"]) async def test_cancellation_before_admission_retains_cleanup_async( self, mock_all_registries: dict[str, Any], phase: str ) -> None: @@ -1581,26 +1622,20 @@ async def test_cancellation_before_admission_retains_cleanup_async( scenario = mock_all_registries["scenario_instance"] run_id = scenario._scenario_result_id record = mock_all_registries["db_result"] + assert isinstance(record, MagicMock) record.scenario_run_state = ScenarioRunState.FAILED memory.get_scenario_result_header_async.return_value = record entered, release, cleanup_entered, cleanup_release = (asyncio.Event() for _ in range(4)) - original_response = service._build_response_async async def prepare_async(*, request: RunScenarioRequest) -> _svc_mod._PreparedRun: record.scenario_run_state = ScenarioRunState.CREATED return _svc_mod._PreparedRun(scenario=scenario) - async def get_results_async(**_: object) -> list[ScenarioResult]: - if phase == "persisted-read": - entered.set() - await release.wait() - return [record] - - async def build_response_async(**kwargs: Any) -> Any: - if phase == "summary-read": + async def get_header_async(**_: object) -> MagicMock: + if phase == "persisted-read" and record.scenario_run_state == ScenarioRunState.CREATED: entered.set() await release.wait() - return await original_response(**kwargs) + return record async def update_state_async(*, scenario_run_state: ScenarioRunState, **_: object) -> None: record.scenario_run_state = scenario_run_state @@ -1619,8 +1654,7 @@ async def cleanup_async(*, expected_states: set[ScenarioRunState], **_: object) service._active_scenario_result_id = "other-run" with ( patch.object(service, "_prepare_run_async", side_effect=prepare_async), - patch.object(memory, "get_scenario_results_async", side_effect=get_results_async), - patch.object(service, "_build_response_async", side_effect=build_response_async), + patch.object(memory, "get_scenario_result_header_async", side_effect=get_header_async), patch.object(memory, "update_scenario_run_state_and_metadata_fields_async", side_effect=update_state_async), patch.object(memory, "try_update_scenario_run_state_async", side_effect=cleanup_async), ): @@ -1700,6 +1734,7 @@ async def _prepare_async(*, request: Any) -> Any: cancelled = _make_db_scenario_result(result_id="cancelled-during-init", run_state=ScenarioRunState.CANCELLED) mock_all_registries["memory"].get_scenario_results_async = AsyncMock(return_value=[cancelled]) + mock_all_registries["memory"].get_scenario_result_header_async = AsyncMock(return_value=cancelled) with patch.object(service, "_prepare_run_async", _prepare_async): with patch.object(service, "_execute_run_async") as execute: response = await service.start_run_async(request=_make_request()) @@ -1740,6 +1775,7 @@ async def test_start_run_honours_a_cancel_that_lands_while_a_failed_run_initiali scenario_instance._scenario_result_id = "cancelled-mid-init" async def _prepare_async(*, request: Any) -> Any: + mock_all_registries["memory"].get_scenario_result_header_async.return_value = cancelled return _svc_mod._PreparedRun(scenario=scenario_instance) cancelled = _make_db_scenario_result(result_id="cancelled-mid-init", run_state=ScenarioRunState.CANCELLED) @@ -1755,8 +1791,9 @@ async def _prepare_async(*, request: Any) -> Any: assert response.status == ScenarioRunState.CANCELLED execute.assert_not_called() - async def test_start_run_does_not_read_a_header_for_a_fresh_run(self, mock_all_registries) -> None: - """A fresh run has no stored state, so it must not pay for an extra query.""" + async def test_start_run_reads_one_header_and_one_returned_summary_for_a_fresh_run( + self, mock_all_registries: dict[str, Any] + ) -> None: service = ScenarioRunService() scenario_instance = mock_all_registries["scenario_instance"] scenario_instance._scenario_result_id = "fresh-id" @@ -1768,7 +1805,9 @@ async def _prepare_async(*, request: Any) -> Any: with patch.object(service, "_execute_run_async"): await service.start_run_async(request=_make_request()) - mock_all_registries["memory"].get_scenario_result_header_async.assert_not_called() + memory = mock_all_registries["memory"] + memory.get_scenario_result_header_async.assert_awaited_once_with(scenario_result_id="fresh-id") + memory.get_scenario_results_async.assert_awaited_once_with(scenario_result_ids=["fresh-id"]) async def test_start_run_failure_does_not_report_a_cancellation(self, mock_all_registries, caplog) -> None: service = ScenarioRunService() @@ -2151,20 +2190,53 @@ async def test_start_run_rejects_missing_result_id(self, mock_all_registries) -> async def test_start_run_rejects_unpersisted_initialized_scenario(self, mock_all_registries) -> None: service = ScenarioRunService() - mock_all_registries["memory"].get_scenario_results_async = AsyncMock(return_value=[]) + mock_all_registries["memory"].get_scenario_result_header_async.return_value = None - with pytest.raises(RuntimeError, match="was not persisted during initialization"): + with ( + patch.object(service, "_enqueue_run_async", new_callable=AsyncMock) as enqueue, + pytest.raises(RuntimeError, match="was not persisted during initialization"), + ): await service.start_run_async(request=_make_request()) - async def test_start_run_cleans_up_when_response_lookup_fails(self, mock_all_registries) -> None: - """A response failure must not leave an active-task entry.""" + enqueue.assert_not_awaited() + mock_all_registries["scenario_instance"].run_async.assert_not_awaited() + mock_all_registries["memory"].get_scenario_results_async.assert_not_awaited() + + async def test_start_run_rejects_changed_resume_id_before_execution_async( + self, mock_all_registries: dict[str, Any] + ) -> None: service = ScenarioRunService() + memory = mock_all_registries["memory"] + memory.get_scenario_result_header_async.return_value = _make_db_scenario_result( + result_id="saved-id", run_state=ScenarioRunState.FAILED + ) + with ( + patch.object(service, "_enqueue_run_async", new_callable=AsyncMock) as enqueue, + pytest.raises(ValueError, match="changed the saved result ID"), + ): + await service.start_run_async(request=_make_request(scenario_result_id="saved-id")) - with patch.object(service, "_build_response_async", return_value=None): - with pytest.raises(RuntimeError, match="not found in the database"): - await service.start_run_async(request=_make_request()) + enqueue.assert_not_awaited() + mock_all_registries["scenario_instance"].run_async.assert_not_awaited() + memory.get_scenario_results_async.assert_not_awaited() - assert service._active_tasks == {} + async def test_start_run_rejects_invalid_header_before_execution_async( + self, mock_all_registries: dict[str, Any] + ) -> None: + service = ScenarioRunService() + mock_all_registries["memory"].get_scenario_result_header_async.side_effect = ValueError( + "Invalid persisted scenario identity" + ) + + with ( + patch.object(service, "_enqueue_run_async", new_callable=AsyncMock) as enqueue, + pytest.raises(ValueError, match="Invalid persisted scenario identity"), + ): + await service.start_run_async(request=_make_request()) + + enqueue.assert_not_awaited() + mock_all_registries["scenario_instance"].run_async.assert_not_awaited() + mock_all_registries["memory"].get_scenario_results_async.assert_not_awaited() async def test_start_run_executes_successfully(self, mock_all_registries) -> None: """The scheduler starts an initialized scenario.""" @@ -2180,6 +2252,44 @@ async def _run_async() -> None: await asyncio.wait_for(released.wait(), timeout=5) await asyncio.sleep(0) + @pytest.mark.parametrize("queued", [False, True]) + async def test_resume_reads_attacks_only_for_returned_summary_async( + self, *, mock_all_registries: dict[str, Any], queued: bool + ) -> None: + attack = AttackResult( + conversation_id="failed-conversation", + objective="test objective", + outcome=AttackOutcome.ERROR, + error_message="429 Too Many Requests", + error_type="RateLimitError", + total_retries=1, + retry_events=[RetryEvent(attempt_number=1, exception_type="RateLimitError", exception_message="429")], + ) + stored = make_scenario_result( + scenario_run_state=ScenarioRunState.FAILED, + attack_results={"direct": [attack]}, + ) + run_id = str(stored.id) + memory = mock_all_registries["memory"] + memory.get_scenario_result_header_async.return_value = stored.model_copy(update={"attack_results": {}}) + memory.get_scenario_results_async.return_value = [stored] + mock_all_registries["scenario_instance"]._scenario_result_id = run_id + service = ScenarioRunService() + if queued: + service._active_scenario_result_id = "other-run" + try: + with patch.object(service, "_execute_run_async", new_callable=AsyncMock): + response = await service.start_run_async(request=_make_request(scenario_result_id=run_id)) + + memory.get_scenario_results_async.assert_awaited_once_with(scenario_result_ids=[run_id]) + assert response.failed_attacks[0].error_message == "429 Too Many Requests" + assert response.total_retries == 1 + assert response.attack_retries[0].retries == attack.retry_events + assert response.queue_position == (1 if queued else None) + finally: + service._active_scenario_result_id = None + await service.shutdown_async() + class TestScenarioRunServiceGetRun: "Tests for ScenarioRunService.get_run_async." @@ -2899,12 +3009,65 @@ async def test_legacy_plan_techniques_agree_across_projections( class TestScenarioRunServiceCancelRun: """Tests for ScenarioRunService.cancel_run_async.""" + @pytest.mark.parametrize("latest_state", [ScenarioRunState.COMPLETED, ScenarioRunState.FAILED]) + async def test_cancel_rechecks_terminal_state_with_headers_async( + self, *, mock_memory: MagicMock, latest_state: ScenarioRunState + ) -> None: + initial = _make_db_scenario_result(run_state=ScenarioRunState.IN_PROGRESS) + latest = _make_db_scenario_result(run_state=latest_state) + mock_memory.get_scenario_result_header_async.side_effect = [initial, latest] + service = ScenarioRunService() + + with pytest.raises(ValueError, match="Cannot cancel run"): + await service.cancel_run_async(scenario_result_id="sr-uuid-1") + + assert mock_memory.get_scenario_result_header_async.await_count == 2 + mock_memory.get_scenario_results_async.assert_not_awaited() + mock_memory.try_update_scenario_run_state_async.assert_not_awaited() + + async def test_cancel_unscheduled_run_reads_details_only_for_response_async(self, mock_memory: MagicMock) -> None: + initial = _make_db_scenario_result(run_state=ScenarioRunState.CREATED) + cancelled = _make_db_scenario_result(run_state=ScenarioRunState.CANCELLED) + mock_memory.get_scenario_result_header_async.return_value = initial + mock_memory.get_scenario_results_async.return_value = [cancelled] + + result = await ScenarioRunService().cancel_run_async(scenario_result_id="sr-uuid-1") + + assert result is not None + assert result.status == ScenarioRunState.CANCELLED + assert mock_memory.get_scenario_result_header_async.await_count == 2 + mock_memory.get_scenario_results_async.assert_awaited_once() + + @pytest.mark.parametrize("skipped_state", [None, ScenarioRunState.CANCELLED, ScenarioRunState.COMPLETED]) + async def test_handoff_uses_headers_to_skip_invalid_queue_entries_async( + self, *, mock_memory: MagicMock, skipped_state: ScenarioRunState | None + ) -> None: + service = ScenarioRunService() + service._active_scenario_result_id = "finished-run" + skipped = _svc_mod._ActiveTask(scenario_result_id="skipped-run") + next_run = _svc_mod._ActiveTask(scenario_result_id="next-run") + service._queued_runs.extend([skipped, next_run]) + mock_memory.get_scenario_result_header_async.side_effect = [ + _make_db_scenario_result(run_state=skipped_state) if skipped_state is not None else None, + _make_db_scenario_result(run_state=ScenarioRunState.QUEUED), + ] + with patch.object(service, "_start_scheduled_run_locked_async", new_callable=AsyncMock) as start: + await service._handoff_scheduler_async(scenario_result_id="finished-run") + + start.assert_awaited_once_with(scheduled=next_run) + assert not service._queued_runs + assert [ + call.kwargs["scenario_result_id"] for call in mock_memory.get_scenario_result_header_async.await_args_list + ] == ["skipped-run", "next-run"] + mock_memory.get_scenario_results_async.assert_not_awaited() + async def test_cancel_run_returns_none_for_unknown_id(self, mock_memory) -> None: """Test that cancel returns None for non-existent run.""" - mock_memory.get_scenario_results_async = AsyncMock(return_value=[]) service = ScenarioRunService() result = await service.cancel_run_async(scenario_result_id="nonexistent-id") assert result is None + mock_memory.get_scenario_result_header_async.assert_awaited_once_with(scenario_result_id="nonexistent-id") + mock_memory.get_scenario_results_async.assert_not_awaited() async def test_cancel_run_sets_cancelled_status(self, mock_all_registries) -> None: """Test that cancelling a running scenario persists CANCELLED to DB.""" @@ -2921,12 +3084,12 @@ async def run_until_cancelled_async() -> None: await asyncio.wait_for(run_started.wait(), timeout=5) # After update_scenario_run_state, the next DB query should return CANCELLED - running_result = mock_all_registries["db_result"] cancelled_result = _make_db_scenario_result( result_id=response.scenario_result_id, run_state=ScenarioRunState.CANCELLED, ) - mock_memory.get_scenario_results_async = AsyncMock(side_effect=[[running_result], [cancelled_result]]) + mock_memory.get_scenario_results_async = AsyncMock(return_value=[cancelled_result]) + mock_memory.get_scenario_result_header_async.reset_mock() result = await service.cancel_run_async(scenario_result_id=response.scenario_result_id) @@ -2939,6 +3102,8 @@ async def run_until_cancelled_async() -> None: ) assert result is not None assert result.status == ScenarioRunState.CANCELLED + mock_memory.get_scenario_result_header_async.assert_awaited_once() + mock_memory.get_scenario_results_async.assert_awaited_once() class TestScenarioRunServiceRecovery: @@ -3062,6 +3227,9 @@ def _update_state(*, scenario_result_id: str, scenario_run_state: ScenarioRunSta mock_scenario_registry.create_and_initialize_async = AsyncMock(side_effect=_create_scenario) mock_memory.get_scenario_results_async = AsyncMock(side_effect=_get_results) + mock_memory.get_scenario_result_header_async = AsyncMock( + side_effect=lambda *, scenario_result_id: records.get(scenario_result_id) + ) mock_memory.update_scenario_run_state_async = AsyncMock(side_effect=_update_state) mock_memory.try_update_scenario_run_state_async = AsyncMock(side_effect=_update_state) mock_memory.update_scenario_run_state_and_metadata_fields_async = AsyncMock(side_effect=_update_state) @@ -3146,6 +3314,9 @@ async def _shutdown_async() -> None: await service.shutdown_async() mock_memory.get_scenario_results_async = AsyncMock(side_effect=_get_results) + mock_memory.get_scenario_result_header_async = AsyncMock( + side_effect=lambda *, scenario_result_id: records.get(scenario_result_id) + ) mock_memory.update_scenario_run_state_async = AsyncMock(side_effect=_update_state) mock_memory.try_update_scenario_run_state_async = AsyncMock(side_effect=_try_update_state) mock_memory.update_scenario_run_state_and_metadata_fields_async = AsyncMock(side_effect=_update_state) @@ -3259,13 +3430,12 @@ async def run_until_cancelled() -> None: response = await service.start_run_async(request=_make_request()) await asyncio.sleep(0) - running_result = mock_all_registries["db_result"] cancelled_result = _make_db_scenario_result( result_id=response.scenario_result_id, run_state=ScenarioRunState.CANCELLED, ) cancelled_result.metadata = {} - mock_memory.get_scenario_results_async = AsyncMock(side_effect=[[running_result], [cancelled_result]]) + mock_memory.get_scenario_results_async = AsyncMock(return_value=[cancelled_result]) await service.cancel_run_async(scenario_result_id=response.scenario_result_id) mock_memory.get_scenario_result_header_async = AsyncMock(return_value=cancelled_result) @@ -3298,20 +3468,12 @@ async def run_until_cancelled() -> None: # Narrowly scoped internal invariant: the live run stays tracked until readback. assert rid in service._active_tasks - running_result = mock_all_registries["db_result"] cancelled_result = _make_db_scenario_result( result_id=rid, run_state=ScenarioRunState.CANCELLED, ) cancelled_result.metadata = {} - mock_memory.get_scenario_results_async = AsyncMock( - side_effect=[ - [running_result], - [cancelled_result], - [cancelled_result], - [cancelled_result], - ] - ) + mock_memory.get_scenario_results_async = AsyncMock(return_value=[cancelled_result]) await service.cancel_run_async(scenario_result_id=rid) fetched = await service.get_run_async(scenario_result_id=rid) @@ -3326,20 +3488,22 @@ async def run_until_cancelled() -> None: async def test_cancel_completed_run_raises_value_error(self, mock_memory) -> None: """Test that cancelling a completed run raises ValueError.""" db_result = _make_db_scenario_result(result_id="sr-done", run_state=ScenarioRunState.COMPLETED) - mock_memory.get_scenario_results_async = AsyncMock(return_value=[db_result]) + mock_memory.get_scenario_result_header_async.return_value = db_result service = ScenarioRunService() with pytest.raises(ValueError, match="Cannot cancel run"): await service.cancel_run_async(scenario_result_id="sr-done") + mock_memory.get_scenario_results_async.assert_not_awaited() async def test_cancel_already_cancelled_run_raises_value_error(self, mock_memory) -> None: """Test that cancelling an already-cancelled run raises ValueError.""" db_result = _make_db_scenario_result(result_id="sr-cancelled", run_state=ScenarioRunState.CANCELLED) - mock_memory.get_scenario_results_async = AsyncMock(return_value=[db_result]) + mock_memory.get_scenario_result_header_async.return_value = db_result service = ScenarioRunService() with pytest.raises(ValueError, match="Cannot cancel run"): await service.cancel_run_async(scenario_result_id="sr-cancelled") + mock_memory.get_scenario_results_async.assert_not_awaited() class TestScenarioRunServiceExecution: