From 1adbc0c7d458a3f28ee9db370977828a1e53caac Mon Sep 17 00:00:00 2001 From: Imran Ahamed Date: Sun, 14 Jun 2026 17:49:44 -0500 Subject: [PATCH] MAINT Migrate internal callers of deprecated ScenarioResult.to_dict()/from_dict() to Pydantic API ScenarioResult.to_dict() and ScenarioResult.from_dict(...) are marked for removal in 0.16.0. Internal PyRIT code still calls the deprecated forms, emitting a DeprecationWarning on every GET /api/scenarios/runs/{id}/results REST call and every CLI scenario result print. Both deprecated methods are thin wrappers around the Pydantic API (to_dict delegates to model_dump(mode='json', by_alias=True); from_dict delegates to model_validate). Output shape is identical, so the change is behavior-preserving. Updates the two production call sites, three docstring references, and the matching test mocks. Also tightens the REST route return annotation from dict to dict[str, Any]. --- pyrit/backend/routes/scenarios.py | 8 +++++--- pyrit/cli/_output.py | 4 ++-- pyrit/cli/api_client.py | 2 +- .../unit/backend/test_scenario_run_routes.py | 2 +- tests/unit/cli/test_output.py | 19 +++++++++++-------- 5 files changed, 20 insertions(+), 15 deletions(-) diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py index 941d8021fb..ce24f4c448 100644 --- a/pyrit/backend/routes/scenarios.py +++ b/pyrit/backend/routes/scenarios.py @@ -12,6 +12,8 @@ /api/scenarios/runs — scenario execution lifecycle """ +from typing import Any + from fastapi import APIRouter, HTTPException, Query, status from pyrit.backend.models.common import ProblemDetail @@ -199,7 +201,7 @@ async def cancel_scenario_run(scenario_result_id: str) -> ScenarioRunSummary: # 409: {"model": ProblemDetail, "description": "Run not yet completed"}, }, ) -async def get_scenario_run_results(scenario_result_id: str) -> dict: # pyrit-async-suffix-exempt +async def get_scenario_run_results(scenario_result_id: str) -> dict[str, Any]: # pyrit-async-suffix-exempt """ Get detailed results for a completed scenario run. @@ -209,7 +211,7 @@ async def get_scenario_run_results(scenario_result_id: str) -> dict: # pyrit-as scenario_result_id: The scenario_result_id. Returns: - dict: ScenarioResult.to_dict() payload. + dict: ``ScenarioResult.model_dump(mode="json", by_alias=True)`` payload. """ service = get_scenario_run_service() try: @@ -222,4 +224,4 @@ async def get_scenario_run_results(scenario_result_id: str) -> dict: # pyrit-as status_code=status.HTTP_404_NOT_FOUND, detail=f"Scenario run '{scenario_result_id}' not found", ) - return result.to_dict() + return result.model_dump(mode="json", by_alias=True) diff --git a/pyrit/cli/_output.py b/pyrit/cli/_output.py index 3580c8e5b6..596fcbd682 100644 --- a/pyrit/cli/_output.py +++ b/pyrit/cli/_output.py @@ -282,12 +282,12 @@ async def print_scenario_result_async(*, result_dict: dict[str, Any]) -> None: Print detailed scenario results using the output module. Args: - result_dict: ``ScenarioResult.to_dict()`` payload from the REST API. + result_dict: ``ScenarioResult.model_dump(mode="json", by_alias=True)`` payload from the REST API. """ from pyrit.models.scenario_result import ScenarioResult from pyrit.output.scenario_result.pretty import PrettyScenarioResultMemoryPrinter - scenario_result = ScenarioResult.from_dict(result_dict) + scenario_result = ScenarioResult.model_validate(result_dict) printer = PrettyScenarioResultMemoryPrinter() await printer.write_async(scenario_result) diff --git a/pyrit/cli/api_client.py b/pyrit/cli/api_client.py index bfd75ca420..937dfad9a4 100644 --- a/pyrit/cli/api_client.py +++ b/pyrit/cli/api_client.py @@ -228,7 +228,7 @@ async def get_scenario_run_results_async(self, *, scenario_result_id: str) -> di Get detailed results for a completed scenario run. Returns: - dict: ``ScenarioResult.to_dict()`` payload. + dict: ``ScenarioResult.model_dump(mode="json", by_alias=True)`` payload. """ return await self._get_json_async(path=f"/api/scenarios/runs/{scenario_result_id}/results") diff --git a/tests/unit/backend/test_scenario_run_routes.py b/tests/unit/backend/test_scenario_run_routes.py index faf40a5b8f..449653a3ce 100644 --- a/tests/unit/backend/test_scenario_run_routes.py +++ b/tests/unit/backend/test_scenario_run_routes.py @@ -232,7 +232,7 @@ class TestGetScenarioRunResultsRoute: def test_get_results_returns_200(self, client: TestClient) -> None: """Test that getting results of a completed run returns 200.""" mock_scenario_result = MagicMock() - mock_scenario_result.to_dict.return_value = { + mock_scenario_result.model_dump.return_value = { "id": "result-uuid", "scenario_identifier": {"name": "foundry.red_team_agent", "version": 1}, "scenario_run_state": "COMPLETED", diff --git a/tests/unit/cli/test_output.py b/tests/unit/cli/test_output.py index 26f41cccc2..e288480fd0 100644 --- a/tests/unit/cli/test_output.py +++ b/tests/unit/cli/test_output.py @@ -319,23 +319,26 @@ async def test_print_scenario_result_async_uses_pretty_printer(): fake_printer.write_async = AsyncMock() with ( - patch("pyrit.models.scenario_result.ScenarioResult.from_dict", return_value=fake_scenario) as from_dict_mock, + patch( + "pyrit.models.scenario_result.ScenarioResult.model_validate", return_value=fake_scenario + ) as model_validate_mock, patch( "pyrit.output.scenario_result.pretty.PrettyScenarioResultMemoryPrinter", return_value=fake_printer ) as printer_cls, ): await _output.print_scenario_result_async(result_dict=result_dict) - from_dict_mock.assert_called_once_with(result_dict) + model_validate_mock.assert_called_once_with(result_dict) printer_cls.assert_called_once_with() fake_printer.write_async.assert_awaited_once_with(fake_scenario) async def test_print_scenario_result_async_roundtrip_with_real_payload(): """ - Integration smoke test: a real ScenarioResult.to_dict() payload must flow - through ScenarioResult.from_dict() inside print_scenario_result_async - without raising. Locks the REST contract used by the CLI thin client. + Integration smoke test: a real ``ScenarioResult.model_dump(mode="json", by_alias=True)`` + payload must flow through ``ScenarioResult.model_validate(...)`` inside + ``print_scenario_result_async`` without raising. Locks the REST contract used by the CLI + thin client. """ from datetime import datetime, timezone @@ -361,10 +364,10 @@ async def test_print_scenario_result_async_roundtrip_with_real_payload(): attack_results={"strat_a": [attack]}, scenario_run_state="COMPLETED", ) - payload = original.to_dict() + payload = original.model_dump(mode="json", by_alias=True) - # Drive print_scenario_result_async through the real from_dict path; only - # stub the printer to keep the test fast. + # Drive print_scenario_result_async through the real model_validate path; + # only stub the printer to keep the test fast. fake_printer = MagicMock() fake_printer.write_async = AsyncMock() with patch(