chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_callback_manager.py

@@ -0,0 +1,86 @@
+import sys
+from unittest.mock import MagicMock
+
+import pytest
+
+from ray.train.v2._internal.execution.callback import ControllerCallback
+from ray.train.v2._internal.execution.callback_manager import CallbackManager
+from ray.train.v2.api.exceptions import ControllerError
+
+
+def test_invoke_callback_without_hook():
+    """Test that callbacks without the hook method raise an error."""
+
+    class CallbackWithoutHook:
+        pass
+
+    manager = CallbackManager([CallbackWithoutHook()])
+    with pytest.raises(ControllerError) as exc_info:
+        manager.invoke("test_hook", "arg1")
+    assert isinstance(exc_info.value.controller_failure, AttributeError)
+
+
+def test_invoke_multiple_callbacks_all_succeed():
+    """Test that multiple callbacks are invoked in sequence."""
+    callback1 = MagicMock()
+    callback1.test_hook = MagicMock()
+    callback2 = MagicMock()
+    callback2.test_hook = MagicMock()
+
+    manager = CallbackManager([callback1, callback2])
+    result = manager.invoke("test_hook", "arg")
+
+    callback1.test_hook.assert_called_once_with("arg")
+    callback2.test_hook.assert_called_once_with("arg")
+    assert result is None
+
+
+def test_invoke_with_real_controller_callback_error_returned():
+    """Test with a real ControllerCallback implementation that raises an error."""
+
+    class MockControllerCallback(ControllerCallback):
+        def __init__(self):
+            self.called = False
+
+        def after_controller_start(self, train_run_context):
+            self.called = True
+            raise ValueError("Intentional error")
+
+    callback = MockControllerCallback()
+
+    manager = CallbackManager([callback])
+    train_run_context = MagicMock()
+    with pytest.raises(ControllerError) as exc_info:
+        manager.invoke("after_controller_start", train_run_context)
+
+    assert callback.called is True
+    assert isinstance(exc_info.value.controller_failure, ValueError)
+
+
+def test_invoke_callback_error_returns_controller_error():
+    callback = MagicMock()
+    callback.test_hook = MagicMock(side_effect=ValueError("Original hook error"))
+
+    manager = CallbackManager([callback])
+    with pytest.raises(ControllerError) as exc_info:
+        manager.invoke("test_hook", "arg1", key1="value1")
+
+    assert isinstance(exc_info.value.controller_failure, ValueError)
+
+
+def test_invoke_best_effort_calls_all_callbacks_even_on_failure():
+    """invoke_best_effort continues calling remaining callbacks after failure."""
+    callback1 = MagicMock()
+    callback1.test_hook = MagicMock(side_effect=ValueError("fail"))
+    callback2 = MagicMock()
+    callback2.test_hook = MagicMock()
+
+    manager = CallbackManager([callback1, callback2])
+    manager.invoke_best_effort("test_hook", "arg")
+
+    callback1.test_hook.assert_called_once_with("arg")
+    callback2.test_hook.assert_called_once_with("arg")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-x", __file__]))