From 3d7bf967e36a643a2949b1aba9262cb57fddfcb1 Mon Sep 17 00:00:00 2001 From: yuqei Date: Thu, 30 Apr 2026 22:41:01 -0400 Subject: [PATCH] fix: scrub training attempt responses --- src/api/training.py | 14 +- .../test_training_public_api_contract_m8.py | 163 ++++++++++++++++++ 2 files changed, 173 insertions(+), 4 deletions(-) diff --git a/src/api/training.py b/src/api/training.py index dbe5a3e..3ea176c 100644 --- a/src/api/training.py +++ b/src/api/training.py @@ -28,13 +28,13 @@ def create_training_attempt(body: dict, service: Optional[TrainingCaseService] = user_id = _get_str(body, "user_id", "local") current_frame_index = _get_int(body, "current_frame_index", 0) meta = _get_dict(body, "meta", {}) - attempt_id = _service(service).create_attempt( + response = _service(service).create_attempt( case_id=case_id, user_id=user_id, current_frame_index=current_frame_index, meta=meta, ) - return sanitize_public_training_payload({"id": attempt_id}) + return _public_id_response(response) def add_training_annotation(attempt_id: int, body: dict, service: Optional[TrainingCaseService] = None) -> dict: @@ -42,14 +42,14 @@ def add_training_annotation(attempt_id: int, body: dict, service: Optional[Train annotation_payload = _require_dict(body, "payload") user_id = _get_str(body, "user_id", "local") task_id = _get_optional_int(body, "task_id") - annotation_id = _service(service).add_annotation( + response = _service(service).add_annotation( attempt_id=attempt_id, annotation_type=annotation_type, payload=annotation_payload, user_id=user_id, task_id=task_id, ) - return sanitize_public_training_payload({"id": annotation_id}) + return _public_id_response(response) def submit_training_attempt(attempt_id: int, body: dict, service: Optional[TrainingCaseService] = None) -> dict: @@ -91,6 +91,12 @@ def _service(service: Optional[TrainingCaseService]) -> TrainingCaseService: return service or TrainingCaseService() +def _public_id_response(response) -> dict: + if isinstance(response, dict): + return sanitize_public_training_payload(response) + return sanitize_public_training_payload({"id": response}) + + def _require_int(body: dict, key: str) -> int: value = _require_value(body, key) try: diff --git a/tests/unit/test_training_public_api_contract_m8.py b/tests/unit/test_training_public_api_contract_m8.py index 77e8afe..e2ae125 100644 --- a/tests/unit/test_training_public_api_contract_m8.py +++ b/tests/unit/test_training_public_api_contract_m8.py @@ -24,8 +24,14 @@ M20_PRIVATE_MARKERS = ( "m20-solution-close", "m20-solution-rubric", ) +M21_PRIVATE_MARKERS = ( + "m21-solution-side", + "m21-solution-close", + "m21-solution-rubric", +) PRIVATE_FIELD_TEXT = "public text mentions " + INTERNAL_FIELD M20_FIELD_TEXT = "m20 public text mentions " + INTERNAL_FIELD +M21_FIELD_TEXT = "m21 public text mentions " + INTERNAL_FIELD DOCUMENTED_TRAINING_ROUTES = ( ("GET", "/training/cases"), @@ -260,6 +266,75 @@ class _M20SubmitService: } +class _M21AttemptLifecycleService: + def __init__(self) -> None: + self.calls = [] + + def create_attempt( + self, + case_id: int, + user_id: str = "local", + current_frame_index: int = 0, + meta=None, + ) -> dict: + self.calls.append( + { + "method": "create_attempt", + "case_id": case_id, + "user_id": user_id, + "current_frame_index": current_frame_index, + "meta": meta, + } + ) + return { + "id": 808, + "status": "started", + "summary": { + "public_result": "attempt created", + "public_note": M21_FIELD_TEXT, + "answer": {"side": M21_PRIVATE_MARKERS[0]}, + INTERNAL_FIELD: {"rubric": M21_PRIVATE_MARKERS[1]}, + "nested": { + "visible": "kept", + "answer": M21_PRIVATE_MARKERS[2], + }, + }, + } + + def add_annotation( + self, + attempt_id: int, + annotation_type: str, + payload: dict, + user_id: str = "local", + task_id=None, + ) -> dict: + self.calls.append( + { + "method": "add_annotation", + "attempt_id": attempt_id, + "annotation_type": annotation_type, + "payload": payload, + "user_id": user_id, + "task_id": task_id, + } + ) + return { + "id": 909, + "status": "recorded", + "summary": { + "public_result": "annotation added", + "public_note": M21_FIELD_TEXT, + "answer": {"side": M21_PRIVATE_MARKERS[0]}, + INTERNAL_FIELD: {"rubric": M21_PRIVATE_MARKERS[1]}, + "nested": { + "visible": "kept", + "answer": M21_PRIVATE_MARKERS[2], + }, + }, + } + + def _documented_routes() -> tuple: text = TRAINING_PRODUCT_DOC.read_text(encoding="utf-8") routes = [] @@ -705,6 +780,94 @@ class TrainingPublicApiContractM8Tests(unittest.TestCase): self.assertNotIn("answer", keys) self.assertNotIn(INTERNAL_FIELD, keys) + def test_public_training_create_attempt_route_scrubs_unexpected_service_solution_material(self) -> None: + service = _M21AttemptLifecycleService() + dependencies = _training_dependencies(service) + + status, payload = server.route_request( + "/training/attempts", + {}, + dependencies=dependencies, + method="POST", + body={"case_id": 17, "user_id": "analyst", "current_frame_index": 2, "meta": {"source": "unit"}}, + ) + + self.assertEqual(status, 201) + self.assertEqual( + service.calls, + [ + { + "method": "create_attempt", + "case_id": 17, + "user_id": "analyst", + "current_frame_index": 2, + "meta": {"source": "unit"}, + } + ], + ) + self.assertEqual(payload["id"], 808) + self.assertEqual(payload["status"], "started") + self.assertEqual(payload["summary"]["public_result"], "attempt created") + self.assertEqual(payload["summary"]["public_note"], "m21 public text mentions private solution field") + self.assertEqual(payload["summary"]["nested"], {"visible": "kept"}) + + serialized = json.dumps(payload, sort_keys=True) + self.assertNotIn(INTERNAL_FIELD, serialized) + for private_text in M21_PRIVATE_MARKERS: + with self.subTest(private_text=private_text): + self.assertNotIn(private_text, serialized) + + keys = _dict_keys(payload) + self.assertNotIn("answer", keys) + self.assertNotIn(INTERNAL_FIELD, keys) + + def test_public_training_annotation_route_scrubs_unexpected_service_solution_material(self) -> None: + service = _M21AttemptLifecycleService() + dependencies = _training_dependencies(service) + + status, payload = server.route_request( + "/training/attempts/808/annotations", + {}, + dependencies=dependencies, + method="POST", + body={ + "annotation_type": "zone", + "payload": {"note": "context"}, + "user_id": "analyst", + "task_id": 33, + }, + ) + + self.assertEqual(status, 201) + self.assertEqual( + service.calls, + [ + { + "method": "add_annotation", + "attempt_id": 808, + "annotation_type": "zone", + "payload": {"note": "context"}, + "user_id": "analyst", + "task_id": 33, + } + ], + ) + self.assertEqual(payload["id"], 909) + self.assertEqual(payload["status"], "recorded") + self.assertEqual(payload["summary"]["public_result"], "annotation added") + self.assertEqual(payload["summary"]["public_note"], "m21 public text mentions private solution field") + self.assertEqual(payload["summary"]["nested"], {"visible": "kept"}) + + serialized = json.dumps(payload, sort_keys=True) + self.assertNotIn(INTERNAL_FIELD, serialized) + for private_text in M21_PRIVATE_MARKERS: + with self.subTest(private_text=private_text): + self.assertNotIn(private_text, serialized) + + keys = _dict_keys(payload) + self.assertNotIn("answer", keys) + self.assertNotIn(INTERNAL_FIELD, keys) + if __name__ == "__main__": unittest.main()