fix tests

Author: Dylan Huang

tests/test_status_migration_changes.py

@@ -83,10 +83,10 @@ def test_trajectory_running_status_creation(self):
         row = EvaluationRow(messages=[])
 
         # Simulate the status assignment from MCP execution manager
-        row.rollout_status = Status(code=Status.Code.OK, message="Rollout is running", details=[])
+        row.rollout_status = Status(code=Status.Code.RUNNING, message="Rollout is running", details=[])
 
         # Verify the status
-        assert row.rollout_status.code == Status.Code.OK
+        assert row.rollout_status.code == Status.Code.RUNNING
         assert row.rollout_status.message == "Rollout is running"
         assert row.rollout_status.is_running()
         assert not row.rollout_status.is_finished()
@@ -278,7 +278,7 @@ def test_status_creation_methods_integration(self):
         # Test running status
         row.rollout_status = Status.rollout_running()
         assert row.rollout_status.is_running()
-        assert row.rollout_status.code == Status.Code.OK
+        assert row.rollout_status.code == Status.Code.RUNNING
 
         # Test finished status
         row.rollout_status = Status.rollout_finished()
@@ -426,7 +426,7 @@ def test_status_code_compliance(self):
 
         # Test standard gRPC codes
         statuses = [
-            (Status.rollout_running(), Status.Code.OK),
+            (Status.rollout_running(), Status.Code.RUNNING),
             (Status.rollout_finished(), Status.Code.FINISHED),  # Custom code
             (Status.rollout_error("Test"), Status.Code.INTERNAL),
         ]

tests/test_status_migration_integration.py

@@ -34,7 +34,7 @@ def test_rollout_status_field_access(self):
         row = EvaluationRow(messages=[])
 
         # Should be able to access directly
-        assert row.rollout_status.code == Status.Code.OK
+        assert row.rollout_status.code == Status.Code.RUNNING
         assert row.rollout_status.message == "Rollout is running"
 
         # Should be able to set directly
@@ -276,7 +276,7 @@ def test_status_code_mapping(self):
 
         # Test different status types and their codes
         statuses = [
-            (Status.rollout_running(), Status.Code.OK),
+            (Status.rollout_running(), Status.Code.RUNNING),
             (Status.rollout_finished(), Status.Code.FINISHED),
             (Status.rollout_error("Test"), Status.Code.INTERNAL),
         ]

tests/test_status_model.py

@@ -82,7 +82,7 @@ def test_status_creation_methods(self):
         """Test the convenience methods for creating Status instances."""
         # Test running status
         running_status = Status.rollout_running()
-        assert running_status.code == Status.Code.OK
+        assert running_status.code == Status.Code.RUNNING
         assert running_status.message == "Rollout is running"
         assert running_status.details == []
 
@@ -267,7 +267,7 @@ def test_evaluation_row_default_status(self):
         assert not hasattr(row, "status")
 
         # Default status should be running
-        assert row.rollout_status.code == Status.Code.OK
+        assert row.rollout_status.code == Status.Code.RUNNING
         assert row.rollout_status.message == "Rollout is running"
         assert row.rollout_status.details == []