fix test_status_migration_integration

Author: Dylan Huang

eval_protocol/models.py

@@ -148,7 +148,7 @@ def rollout_error(cls, error_message: str, extra_info: Optional[Dict[str, Any]]
         """Create a status indicating the rollout failed with an error."""
         details = []
         if extra_info:
-            details.append(ErrorInfo.rollout_error(extra_info).to_aip193_format())
+            details.append(ErrorInfo.extra_info(extra_info).to_aip193_format())
         return cls.error(error_message, details)
 
     @classmethod
@@ -210,7 +210,7 @@ def get_extra_info(self) -> Optional[Dict[str, Any]]:
             metadata = detail.get("metadata", {})
             reason = detail.get("reason")
             # Skip termination_reason and stopped details, return other error info
-            if reason not in [ErrorInfo.REASON_TERMINATION_REASON, ErrorInfo.REASON_STOPPED]:
+            if reason in [ErrorInfo.REASON_EXTRA_INFO]:
                 return metadata
         return None
 

tests/test_status_migration_integration.py

@@ -106,33 +106,33 @@ def test_termination_reason_in_status_details(self):
         row = EvaluationRow(messages=[])
 
         # Set status with termination reason
-        termination_status = Status.with_termination_reason("goal_reached")
+        termination_status = Status.with_termination_reason(TerminationReason.CONTROL_PLANE_SIGNAL)
         row.rollout_status = termination_status
 
         # Should be finished
         assert row.rollout_status.is_finished()
 
         # Should have termination reason in details
-        assert row.rollout_status.get_termination_reason() == "goal_reached"
+        assert row.rollout_status.get_termination_reason() == TerminationReason.CONTROL_PLANE_SIGNAL
 
         # Check details structure
         assert len(row.rollout_status.details) == 1
         detail = row.rollout_status.details[0]
         assert detail["@type"] == "type.googleapis.com/google.rpc.ErrorInfo"
         assert detail["reason"] == "TERMINATION_REASON"
         assert detail["domain"] == "evalprotocol.io"
-        assert detail["metadata"]["termination_reason"] == "goal_reached"
+        assert detail["metadata"]["termination_reason"] == TerminationReason.CONTROL_PLANE_SIGNAL
 
     def test_termination_reason_with_extra_info(self):
         """Test termination reason with additional extra info."""
         row = EvaluationRow(messages=[])
 
         extra_info = {"steps": 10, "reward": 0.8}
-        termination_status = Status.with_termination_reason("timeout", extra_info)
+        termination_status = Status.with_termination_reason(TerminationReason.USER_STOP, extra_info)
         row.rollout_status = termination_status
 
         # Should have both termination reason and extra info
-        assert row.rollout_status.get_termination_reason() == "timeout"
+        assert row.rollout_status.get_termination_reason() == TerminationReason.USER_STOP
         assert row.rollout_status.get_extra_info() == extra_info
 
         # Check details structure
@@ -157,21 +157,21 @@ def test_multiple_termination_reasons(self):
                 "@type": "type.googleapis.com/google.rpc.ErrorInfo",
                 "reason": "TERMINATION_REASON",
                 "domain": "evalprotocol.io",
-                "metadata": {"termination_reason": "first"},
+                "metadata": {"termination_reason": TerminationReason.USER_STOP},
             },
             {
                 "@type": "type.googleapis.com/google.rpc.ErrorInfo",
                 "reason": "TERMINATION_REASON",
                 "domain": "evalprotocol.io",
-                "metadata": {"termination_reason": "second"},
+                "metadata": {"termination_reason": TerminationReason.SKIPPABLE_ERROR},
             },
         ]
 
         status = Status(code=Status.Code.FINISHED, message="Test", details=details)
         row.rollout_status = status
 
         # Should return the first termination reason found
-        assert row.rollout_status.get_termination_reason() == "first"
+        assert row.rollout_status.get_termination_reason() == TerminationReason.USER_STOP
 
 
 class TestErrorHandlingIntegration:
@@ -204,7 +204,7 @@ def test_error_status_with_metadata(self):
         assert len(row.rollout_status.details) == 1
         detail = row.rollout_status.details[0]
         assert detail["@type"] == "type.googleapis.com/google.rpc.ErrorInfo"
-        assert detail["reason"] == "ROLLOUT_ERROR"
+        assert detail["reason"] == "EXTRA_INFO"
         assert detail["domain"] == "evalprotocol.io"
         assert detail["metadata"] == error_info
 
@@ -331,7 +331,7 @@ def test_status_model_validate(self):
 
         # Set a complex status
         extra_info = {"steps": 10, "reward": 0.8}
-        original_status = Status.with_termination_reason("goal_reached", extra_info)
+        original_status = Status.with_termination_reason(TerminationReason.CONTROL_PLANE_SIGNAL, extra_info)
         row.rollout_status = original_status
 
         # Dump and reconstruct
@@ -344,7 +344,7 @@ def test_status_model_validate(self):
         assert len(reconstructed_status.details) == len(original_status.details)
 
         # Should preserve functionality
-        assert reconstructed_status.get_termination_reason() == "goal_reached"
+        assert reconstructed_status.get_termination_reason() == TerminationReason.CONTROL_PLANE_SIGNAL
         assert reconstructed_status.get_extra_info() == extra_info
 
 
@@ -377,7 +377,7 @@ def test_malformed_status_details(self):
         row.rollout_status = malformed_status
 
         # Should handle gracefully
-        assert row.rollout_status.get_termination_reason() == "test"
+        assert row.rollout_status.get_termination_reason() is None
         assert row.rollout_status.get_extra_info() is None
 
     def test_large_metadata_handling(self):