fix test

Author: xzrderek

eval_protocol/exceptions.py

@@ -160,17 +160,18 @@ class ScoreInvalidError(EvalProtocolError):
 }
 
 
-def exception_for_status_code(code: int) -> Optional[EvalProtocolError]:
+def exception_for_status_code(code: int, message: str = "") -> Optional[EvalProtocolError]:
     """
     Create an exception instance for a given status code.
 
     Args:
         code: Status code from Status.Code enum
+        message: Optional error message to include in the exception
 
     Returns:
         Exception instance or None if code is OK (0)
     """
     exception_class = STATUS_CODE_TO_EXCEPTION.get(code)
     if exception_class is None:
         return None
-    return exception_class()
+    return exception_class(message) if message else exception_class()

eval_protocol/pytest/remote_rollout_processor.py

@@ -161,8 +161,8 @@ def _get_status() -> Dict[str, Any]:
                         f"Found Fireworks log for rollout {row.execution_metadata.rollout_id} with status code {status_code}"
                     )
 
-                    # Create and raise exception if appropriate
-                    exception = exception_for_status_code(status_code)
+                    # Create and raise exception if appropriate, preserving original message
+                    exception = exception_for_status_code(status_code, status_message)
                     if exception is not None:
                         raise exception
 

tests/test_exceptions.py

@@ -348,3 +348,24 @@ def test_integration_with_retry_logic():
         assert exception_class in DEFAULT_RETRYABLE_EXCEPTIONS, (
             f"{exception_class.__name__} should be in DEFAULT_RETRYABLE_EXCEPTIONS for retry support"
         )
+
+
+def test_exception_message_preservation():
+    """Test that error messages are properly preserved in exceptions."""
+    test_cases = [
+        (13, "test error", InternalError),
+        (5, "Model xyz not found", NotFoundError),
+        (7, "Invalid API key", PermissionDeniedError),
+    ]
+
+    for status_code, message, expected_exception_class in test_cases:
+        # Test with message
+        exception = exception_for_status_code(status_code, message)
+        assert exception is not None
+        assert isinstance(exception, expected_exception_class)
+        assert str(exception) == message, f"Exception should preserve message '{message}'"
+
+        # Test without message (should still work)
+        exception_no_msg = exception_for_status_code(status_code)
+        assert exception_no_msg is not None
+        assert isinstance(exception_no_msg, expected_exception_class)