wehub-resource-sync
87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
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__]))
|