def test_process_wrapper_successful_call(
self, mock_manager_class, mock_import, mock_register, mock_get_logger,
mock_initialize_logging):
"""Tests _process_wrapper for a process where there are no errors."""
mock_manager = mock_manager_class.return_value
mock_logger = mock_get_logger.return_value
multiprocessing_queue = multiprocessing_utils.get_context().Queue()
return_queue = mock.MagicMock(spec=multiprocessing_queue)
error_queue = mock.MagicMock(spec=multiprocessing_queue)
logging_queue = mock.MagicMock(spec=multiprocessing_queue)
process_id = "1"
mock_function = mock.MagicMock()
mock_function.__name__ = "mock_function"
args = (1, 2)
kwargs = {"foo": "bar"}
parallel_utils._process_wrapper(
return_queue=return_queue,
error_queue=error_queue,
logging_queue=logging_queue,
process_id=process_id,
extension_package_import_paths=["foo.package", "bar.package"],
call_spec=parallel_utils.CallSpec(mock_function, *args, **kwargs))
mock_initialize_logging.assert_called_once_with(logging_queue)
mock_get_logger.assert_called_once()
mock_logger.debug.assert_called()
mock_import.assert_has_calls(
[mock.call("foo.package"), mock.call("bar.package")])
self.assertEqual(mock_register.call_count, 2)
mock_manager_class.assert_called_once()
mock_function.assert_called_once_with(mock_manager, *args, **kwargs)
return_queue.put.assert_called_once_with(
(process_id, mock_function.return_value))
error_queue.put.assert_not_called()
mock_manager.close.assert_called_once()
def test_process_wrapper_exception_call(
self, mock_manager_class, mock_import, mock_register, mock_get_logger,
mock_initialize_logging):
"""Tests _process_wrapper for a process where function raises an error."""
mock_manager = mock_manager_class.return_value
mock_logger = mock_get_logger.return_value
multiprocessing_queue = multiprocessing_utils.get_context().Queue()
return_queue = mock.MagicMock(spec=multiprocessing_queue)
error_queue = mock.MagicMock(spec=multiprocessing_queue)
process_id = "1"
mock_function = mock.MagicMock()
mock_function.__name__ = "mock_function"
mock_function.side_effect = RuntimeError("Something went wrong")
args = (1, 2)
kwargs = {"foo": "bar"}
parallel_utils._process_wrapper(
return_queue=return_queue,
error_queue=error_queue,
logging_queue=mock.MagicMock(spec=multiprocessing_queue),
process_id=process_id,
# "foo.package" imports but fails registration.
# "bar.package" fails to import.
extension_package_import_paths=["foo.package", "bar.package"],
call_spec=parallel_utils.CallSpec(mock_function, *args, **kwargs))
mock_import.assert_has_calls(
[mock.call("foo.package"), mock.call("bar.package")])
mock_register.assert_called_once()
mock_manager_class.assert_called_once()
mock_function.assert_called_once_with(mock_manager, *args, **kwargs)
mock_logger.warning.assert_called()
return_queue.put.assert_not_called()
error_queue.put.assert_called_once_with(
(process_id, (RuntimeError.__name__, "Something went wrong")))
mock_manager.close.assert_called_once()
def _issue_devices(self, devices: Sequence[str], attribute_name: str,
timeout: float, method_args: Tuple[Any, ...],
method_kwargs: Mapping[str, Any]) -> List[Any]:
"""Executes a device method in parallel on multiple devices.
Args:
devices: Device identifiers to execute the method on.
attribute_name: Name of device method or property to execute in parallel.
timeout: Maximum amount of seconds to allow parallel methods to complete.
method_args: Positional arguments to the method.
method_kwargs: Keyword arguments to the method.
Returns:
Results from the parallel method calls.
"""
call_specs = [
parallel_utils.CallSpec(_run_device_action, device, attribute_name,
method_args, method_kwargs)
for device in devices
]
results, _ = parallel_utils.execute_concurrently(
call_specs=call_specs,
timeout=timeout,
raise_on_process_error=True)
return results
def test_issue_devices_executes_device_method(self, devices_spec,
mock_execute):
"""Tests issue_devices_* executing device methods in parallel."""
method = "shell"
method_args = ("whoami", )
method_kwargs = {"port": 1}
expected_method_return_values = ["user 1", "user 2"]
mock_execute.return_value = (expected_method_return_values,
[None, None])
if isinstance(devices_spec, str):
method_return_values = self.uut.issue_devices_match(
devices_spec, method, *method_args, **method_kwargs)
elif devices_spec is None:
method_return_values = self.uut.issue_devices_all(
method, *method_args, **method_kwargs)
else:
method_return_values = self.uut.issue_devices(
devices_spec, method, *method_args, **method_kwargs)
self.assertEqual(method_return_values, expected_method_return_values)
expected_device_ids = ["sshdevice-0000", "sshdevice-0001"]
expected_call_specs = [
parallel_utils.CallSpec(fire_manager._run_device_action, device,
method, method_args, method_kwargs)
for device in expected_device_ids
]
mock_execute.assert_called_once_with(call_specs=expected_call_specs,
timeout=mock.ANY,
raise_on_process_error=True)
_TEST_EXCEPTION = RuntimeError("Something went wrong.")
def _test_function_raises_exception(manager_inst: manager.Manager) -> NoReturn:
"""Function which raises an error."""
raise _TEST_EXCEPTION
def _test_function_times_out(manager_inst: manager.Manager) -> None:
"""Function which is designed to time out."""
assert isinstance(manager_inst, manager.Manager)
time.sleep(15)
_GOOD_CALL_SPECS = [
parallel_utils.CallSpec(_test_function_with_return, 5),
parallel_utils.CallSpec(_test_function_with_return, some_arg=10),
parallel_utils.CallSpec(_test_function_no_return),
]
_GOOD_CALL_RESULTS = [5, 10, None]
_GOOD_CALL_ERRORS = [None, None, None]
_BAD_CALL_SPECS = [
parallel_utils.CallSpec(_test_function_raises_exception),
parallel_utils.CallSpec(_test_function_times_out),
]
_BAD_CALL_RESULTS = [parallel_utils.NO_RESULT] * len(_BAD_CALL_SPECS)
_BAD_CALL_ERRORS = [
(type(_TEST_EXCEPTION).__name__, str(_TEST_EXCEPTION)),
(errors.ResultNotReceivedError.__name__,
"Did not receive any results from the process."),
]