From 5e46ebd04daa6dccdf5e874d3cf36a04e7d05d28 Mon Sep 17 00:00:00 2001 From: Gabriel Levcovitz Date: Mon, 22 Jul 2024 18:05:46 -0300 Subject: [PATCH] refactor(side-dag): refactor node processes management --- Makefile | 12 +- extras/custom_tests.sh | 40 +++ extras/custom_tests/__init__.py | 0 extras/custom_tests/side_dag/__init__.py | 0 .../custom_tests/side_dag/test_both_fail.py | 98 ++++++++ .../custom_tests/side_dag/test_one_fails.py | 97 +++++++ extras/custom_tests/side_dag/utils.py | 78 ++++++ hathor/cli/side_dag.py | 237 +++++++++--------- tests/cli/test_side_dag.py | 141 +---------- 9 files changed, 443 insertions(+), 260 deletions(-) create mode 100644 extras/custom_tests.sh create mode 100644 extras/custom_tests/__init__.py create mode 100644 extras/custom_tests/side_dag/__init__.py create mode 100644 extras/custom_tests/side_dag/test_both_fail.py create mode 100644 extras/custom_tests/side_dag/test_one_fails.py create mode 100644 extras/custom_tests/side_dag/utils.py diff --git a/Makefile b/Makefile index c6ea165726..eed3cee4ec 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -py_sources = hathor/ tests/ +py_sources = hathor/ tests/ extras/custom_tests/ .PHONY: all all: check tests @@ -49,8 +49,12 @@ tests-genesis: tests-ci: pytest $(tests_ci) +.PHONY: tests-custom +tests-custom: + bash ./extras/custom_tests.sh + .PHONY: tests -tests: tests-cli tests-lib tests-genesis tests-ci +tests: tests-cli tests-lib tests-genesis tests-custom tests-ci .PHONY: tests-full tests-full: @@ -60,11 +64,11 @@ tests-full: .PHONY: mypy mypy: - mypy -p hathor -p tests + mypy -p hathor -p tests -p extras.custom_tests .PHONY: dmypy dmypy: - dmypy run --timeout 86400 -- -p hathor -p tests + dmypy run --timeout 86400 -- -p hathor -p tests -p extras.custom_tests .PHONY: flake8 flake8: diff --git a/extras/custom_tests.sh b/extras/custom_tests.sh new file mode 100644 index 0000000000..ec548c25e9 --- /dev/null +++ b/extras/custom_tests.sh @@ -0,0 +1,40 @@ +#!/bin/bash + +# Define colors +RED='\033[0;31m' +GREEN='\033[0;32m' +BLUE='\033[0;34m' +NC='\033[0m' # No Color + +TESTS_DIR="extras/custom_tests" + +# List of test scripts to be executed +tests=( + /side_dag/test_one_fails.py + /side_dag/test_both_fail.py +) + +# Initialize a variable to track if any test fails +any_test_failed=0 + +# Loop over all tests +for test in "${tests[@]}"; do + echo -e "${BLUE}Testing $test${NC}" + PYTHONPATH=$TESTS_DIR python $TESTS_DIR/$test + result=$? + if [ $result -ne 0 ]; then + echo -e "${RED}Test $test FAILED${NC}" + any_test_failed=1 + else + echo -e "${GREEN}Test $test PASSED${NC}" + fi +done + +# Exit with code 0 if no test failed, otherwise exit with code 1 +if [ $any_test_failed -eq 0 ]; then + echo -e "${GREEN}All tests PASSED${NC}" + exit 0 +else + echo -e "${RED}Some tests FAILED${NC}" + exit 1 +fi diff --git a/extras/custom_tests/__init__.py b/extras/custom_tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/extras/custom_tests/side_dag/__init__.py b/extras/custom_tests/side_dag/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/extras/custom_tests/side_dag/test_both_fail.py b/extras/custom_tests/side_dag/test_both_fail.py new file mode 100644 index 0000000000..a715f744ac --- /dev/null +++ b/extras/custom_tests/side_dag/test_both_fail.py @@ -0,0 +1,98 @@ +# Copyright 2024 Hathor Labs +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import shlex +import signal +import subprocess +import sys + +from utils import ( # type: ignore[import-not-found] + COMMAND, + HATHOR_PROCESS_NAME, + KILL_WAIT_DELAY, + MONITOR_PROCESS_NAME, + SIDE_DAG_PROCESS_NAME, + get_pid_by_name, + is_alive, + popen_is_alive, + wait_seconds, +) + +if sys.platform == 'win32': + print("test skipped on windows") + sys.exit(0) + + +def test_both_fail() -> None: + # Assert that there are no existing processes + assert get_pid_by_name(MONITOR_PROCESS_NAME) is None + assert get_pid_by_name(HATHOR_PROCESS_NAME) is None + assert get_pid_by_name(SIDE_DAG_PROCESS_NAME) is None + + # Run the python command + args = shlex.split(COMMAND) + monitor_process = subprocess.Popen(args) + print(f'running "run_node_with_side_dag" in the background with pid: {monitor_process.pid}') + print('awaiting subprocesses initialization...') + wait_seconds(5) + + monitor_process_pid = get_pid_by_name(MONITOR_PROCESS_NAME) + hathor_process_pid = get_pid_by_name(HATHOR_PROCESS_NAME) + side_dag_process_pid = get_pid_by_name(SIDE_DAG_PROCESS_NAME) + + # Assert that the processes exist and are alive + assert monitor_process_pid == monitor_process.pid + assert monitor_process_pid is not None + assert hathor_process_pid is not None + assert side_dag_process_pid is not None + + assert is_alive(monitor_process_pid) + assert is_alive(hathor_process_pid) + assert is_alive(side_dag_process_pid) + + print('processes are running:') + print(f' "{MONITOR_PROCESS_NAME}" pid: {monitor_process_pid}') + print(f' "{HATHOR_PROCESS_NAME}" pid: {hathor_process_pid}') + print(f' "{SIDE_DAG_PROCESS_NAME}" pid: {side_dag_process_pid}') + print('letting processes run for a while...') + wait_seconds(10) + + # Terminate both subprocess + print('terminating subprocesses...') + os.kill(hathor_process_pid, signal.SIGTERM) + os.kill(side_dag_process_pid, signal.SIGTERM) + print('awaiting processes termination...') + wait_seconds(KILL_WAIT_DELAY, break_function=lambda: not popen_is_alive(monitor_process)) + + # Assert that all process are terminated + assert not popen_is_alive(monitor_process) + assert not is_alive(monitor_process_pid) + assert not is_alive(hathor_process_pid) + assert not is_alive(side_dag_process_pid) + + print('all processes are dead. test succeeded!') + + +try: + test_both_fail() +except Exception: + if monitor_process_pid := get_pid_by_name(MONITOR_PROCESS_NAME): + os.kill(monitor_process_pid, signal.SIGKILL) + if hathor_process_pid := get_pid_by_name(HATHOR_PROCESS_NAME): + os.kill(hathor_process_pid, signal.SIGKILL) + if side_dag_process_pid := get_pid_by_name(SIDE_DAG_PROCESS_NAME): + os.kill(side_dag_process_pid, signal.SIGKILL) + + raise diff --git a/extras/custom_tests/side_dag/test_one_fails.py b/extras/custom_tests/side_dag/test_one_fails.py new file mode 100644 index 0000000000..436b871796 --- /dev/null +++ b/extras/custom_tests/side_dag/test_one_fails.py @@ -0,0 +1,97 @@ +# Copyright 2024 Hathor Labs +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import shlex +import signal +import subprocess +import sys + +from utils import ( # type: ignore[import-not-found] + COMMAND, + HATHOR_PROCESS_NAME, + KILL_WAIT_DELAY, + MONITOR_PROCESS_NAME, + SIDE_DAG_PROCESS_NAME, + get_pid_by_name, + is_alive, + popen_is_alive, + wait_seconds, +) + +if sys.platform == 'win32': + print("test skipped on windows") + sys.exit(0) + + +def test_one_fails() -> None: + # Assert that there are no existing processes + assert get_pid_by_name(MONITOR_PROCESS_NAME) is None + assert get_pid_by_name(HATHOR_PROCESS_NAME) is None + assert get_pid_by_name(SIDE_DAG_PROCESS_NAME) is None + + # Run the python command + args = shlex.split(COMMAND) + monitor_process = subprocess.Popen(args) + print(f'running "run_node_with_side_dag" in the background with pid: {monitor_process.pid}') + print('awaiting subprocesses initialization...') + wait_seconds(5) + + monitor_process_pid = get_pid_by_name(MONITOR_PROCESS_NAME) + hathor_process_pid = get_pid_by_name(HATHOR_PROCESS_NAME) + side_dag_process_pid = get_pid_by_name(SIDE_DAG_PROCESS_NAME) + + # Assert that the processes exist and are alive + assert monitor_process_pid == monitor_process.pid + assert monitor_process_pid is not None + assert hathor_process_pid is not None + assert side_dag_process_pid is not None + + assert is_alive(monitor_process_pid) + assert is_alive(hathor_process_pid) + assert is_alive(side_dag_process_pid) + + print('processes are running:') + print(f' "{MONITOR_PROCESS_NAME}" pid: {monitor_process_pid}') + print(f' "{HATHOR_PROCESS_NAME}" pid: {hathor_process_pid}') + print(f' "{SIDE_DAG_PROCESS_NAME}" pid: {side_dag_process_pid}') + print('letting processes run for a while...') + wait_seconds(10) + + # Terminate one subprocess + print('terminating side-dag process...') + os.kill(side_dag_process_pid, signal.SIGTERM) + print('awaiting process termination...') + wait_seconds(KILL_WAIT_DELAY, break_function=lambda: not popen_is_alive(monitor_process)) + + # Assert that all process are terminated + assert not popen_is_alive(monitor_process) + assert not is_alive(monitor_process_pid) + assert not is_alive(hathor_process_pid) + assert not is_alive(side_dag_process_pid) + + print('all processes are dead. test succeeded!') + + +try: + test_one_fails() +except Exception: + if monitor_process_pid := get_pid_by_name(MONITOR_PROCESS_NAME): + os.kill(monitor_process_pid, signal.SIGKILL) + if hathor_process_pid := get_pid_by_name(HATHOR_PROCESS_NAME): + os.kill(hathor_process_pid, signal.SIGKILL) + if side_dag_process_pid := get_pid_by_name(SIDE_DAG_PROCESS_NAME): + os.kill(side_dag_process_pid, signal.SIGKILL) + + raise diff --git a/extras/custom_tests/side_dag/utils.py b/extras/custom_tests/side_dag/utils.py new file mode 100644 index 0000000000..7793c30acc --- /dev/null +++ b/extras/custom_tests/side_dag/utils.py @@ -0,0 +1,78 @@ +# Copyright 2024 Hathor Labs +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import subprocess +import time +from typing import Callable + +MONITOR_PROCESS_NAME = 'hathor-core: monitor' +PROCNAME_SUFFIX = 'hathor-core' +HATHOR_PROCESS_PREFIX = 'hathor:' +SIDE_DAG_PROCESS_PREFIX = 'side-dag:' +HATHOR_PROCESS_NAME = HATHOR_PROCESS_PREFIX + PROCNAME_SUFFIX +SIDE_DAG_PROCESS_NAME = SIDE_DAG_PROCESS_PREFIX + PROCNAME_SUFFIX +KILL_WAIT_DELAY = 305 + +COMMAND = f""" + python -m hathor run_node_with_side_dag + --disable-logs + --testnet + --memory-storage + --x-localhost-only + --procname-prefix {HATHOR_PROCESS_PREFIX} + --side-dag-testnet + --side-dag-memory-storage + --side-dag-x-localhost-only + --side-dag-procname-prefix {SIDE_DAG_PROCESS_PREFIX} +""" + + +def wait_seconds(seconds: int, *, break_function: Callable[[], bool] | None = None) -> None: + while seconds > 0: + print(f'waiting {seconds} seconds...') + time.sleep(1) + seconds -= 1 + if break_function and break_function(): + break + + +def get_pid_by_name(process_name: str) -> int | None: + try: + output = subprocess.check_output(['pgrep', '-f', process_name], text=True) + except subprocess.CalledProcessError: + return None + pids = output.strip().split() + assert len(pids) <= 1 + try: + return int(pids[0]) + except IndexError: + return None + + +def is_alive(pid: int) -> bool: + try: + os.kill(pid, 0) + except OSError: + return False + return True + + +def popen_is_alive(popen: subprocess.Popen) -> bool: + try: + popen.wait(0) + except subprocess.TimeoutExpired: + return True + assert popen.returncode is not None + return False diff --git a/hathor/cli/side_dag.py b/hathor/cli/side_dag.py index 7bdc270a01..20132f2d53 100644 --- a/hathor/cli/side_dag.py +++ b/hathor/cli/side_dag.py @@ -17,13 +17,14 @@ import os import signal import sys -import traceback +import time from argparse import ArgumentParser from dataclasses import dataclass from enum import Enum -from multiprocessing import Pipe, Process +from multiprocessing import Process from typing import TYPE_CHECKING, Any +from setproctitle import setproctitle from typing_extensions import assert_never, override from hathor.cli.run_node_args import RunNodeArgs # skip-cli-import-custom-check @@ -31,13 +32,6 @@ if TYPE_CHECKING: from hathor.cli.util import LoggingOutput - # Workaround for a typing issue in Windows - if sys.platform == 'win32': - from multiprocessing.connection import _ConnectionBase as Connection - else: - from multiprocessing.connection import Connection - -import psutil from structlog import get_logger from hathor.cli.run_node import RunNode @@ -45,31 +39,16 @@ logger = get_logger() PRE_SETUP_LOGGING: bool = False -HATHOR_NODE_INIT_WAIT_PERIOD: int = 10 - - -class SideDagArgs(RunNodeArgs): - poa_signer_file: str | None - -@dataclass(frozen=True, slots=True) -class HathorProcessInitFail: - reason: str +# Period in seconds for polling subprocesses' state. +MONITOR_WAIT_PERIOD: int = 10 +# Delay in seconds before killing a subprocess after trying to terminate it. +KILL_WAIT_DELAY = 300 -@dataclass(frozen=True, slots=True) -class HathorProcessInitSuccess: - pass - -@dataclass(frozen=True, slots=True) -class HathorProcessTerminated: - pass - - -@dataclass(frozen=True, slots=True) -class SideDagProcessTerminated: - pass +class SideDagArgs(RunNodeArgs): + poa_signer_file: str | None class SideDagRunNode(RunNode): @@ -116,21 +95,39 @@ def main(capture_stdout: bool) -> None: argv = sys.argv[1:] hathor_logging_output, side_dag_logging_output = _process_logging_output(argv) hathor_node_argv, side_dag_argv = _partition_argv(argv) - conn1, conn2 = Pipe() - hathor_node_process = _start_hathor_node_process( - hathor_node_argv, logging_output=hathor_logging_output, capture_stdout=capture_stdout, conn=conn1 - ) - log_options = process_logging_options(side_dag_argv) + # the main process uses the same configuration as the hathor process + log_options = process_logging_options(hathor_node_argv.copy()) setup_logging( - logging_output=side_dag_logging_output, + logging_output=hathor_logging_output, logging_options=log_options, capture_stdout=capture_stdout, - extra_log_info=_get_extra_log_info('side-dag') + extra_log_info=_get_extra_log_info('monitor') + ) + + hathor_node_process = _start_node_process( + argv=hathor_node_argv, + runner=RunNode, + logging_output=hathor_logging_output, + capture_stdout=capture_stdout, + name='hathor', + ) + side_dag_node_process = _start_node_process( + argv=side_dag_argv, + runner=SideDagRunNode, + logging_output=side_dag_logging_output, + capture_stdout=capture_stdout, + name='side-dag', + ) + + logger.info( + 'starting nodes', + monitor_pid=os.getpid(), + hathor_node_pid=hathor_node_process.pid, + side_dag_node_pid=side_dag_node_process.pid ) - logger.info('starting nodes', hathor_node_pid=hathor_node_process.pid, side_dag_pid=os.getpid()) - _run_side_dag_node(side_dag_argv, hathor_node_process=hathor_node_process, conn=conn2) + _run_monitor_process(hathor_node_process, side_dag_node_process) def _process_logging_output(argv: list[str]) -> tuple[LoggingOutput, LoggingOutput]: @@ -205,100 +202,106 @@ def is_option(arg_: str) -> bool: return hathor_node_argv, side_dag_argv -def _run_side_dag_node(argv: list[str], *, hathor_node_process: Process, conn: 'Connection') -> None: +def _run_monitor_process(*processes: Process) -> None: """Function to be called by the main process to run the side-dag full node.""" - logger.info('waiting for hathor node to initialize...') - while not conn.poll(HATHOR_NODE_INIT_WAIT_PERIOD): - logger.info('still waiting for hathor node to initialize...') - - message = conn.recv() - if isinstance(message, HathorProcessInitFail): - logger.critical(f'side-dag node not started because hathor node initialization failed:\n{message.reason}') - return - - assert isinstance(message, HathorProcessInitSuccess) - logger.info('hathor node initialized') - logger.info('starting side-dag node...') - - try: - side_dag = SideDagRunNode(argv=argv) - except (BaseException, Exception): - logger.exception('terminating hathor node due to exception in side-dag node...') - conn.send(SideDagProcessTerminated()) - hathor_node_process.terminate() - return - - side_dag.run() - - # If `run()` returns, either the hathor node exited and terminated us, leaving the message below, or the side-dag - # node exited and we will terminate the hathor node. - if conn.poll(): - message = conn.recv() - assert isinstance(message, HathorProcessTerminated) - logger.critical('side-dag node terminated because hathor node exited') - return - - conn.send(SideDagProcessTerminated()) - logger.critical('terminating hathor node...') - hathor_node_process.terminate() - - -def _start_hathor_node_process( - argv: list[str], + setproctitle('hathor-core: monitor') + signal.signal(signal.SIGINT, lambda _, __: _terminate_and_exit(processes, reason='received SIGINT')) + signal.signal(signal.SIGTERM, lambda _, __: _terminate_and_exit(processes, reason='received SIGTERM')) + + while True: + time.sleep(MONITOR_WAIT_PERIOD) + for process in processes: + if not process.is_alive(): + _terminate_and_exit( + processes, + reason=f'process "{process.name}" (pid: {process.pid}) exited' + ) + + +def _terminate_and_exit(processes: tuple[Process, ...], *, reason: str) -> None: + """Terminate all processes that are alive. Kills them if they're not terminated after a while. + Then, exits the program.""" + logger.critical(f'terminating all nodes. reason: {reason}') + + for process in processes: + if process.is_alive(): + logger.critical(f'terminating process "{process.name}" (pid: {process.pid})...') + process.terminate() + + now = time.time() + while True: + time.sleep(MONITOR_WAIT_PERIOD) + if time.time() >= now + KILL_WAIT_DELAY: + _kill_all(processes) + break + + all_are_dead = all(not process.is_alive() for process in processes) + if all_are_dead: + break + + logger.critical('all nodes terminated.') + sys.exit(0) + + +def _kill_all(processes: tuple[Process, ...]) -> None: + """Kill all processes that are alive.""" + for process in processes: + if process.is_alive(): + logger.critical(f'process "{process.name}" (pid: {process.pid}) still alive, killing it...') + process.kill() + + +def _start_node_process( *, + argv: list[str], + runner: type[RunNode], logging_output: LoggingOutput, capture_stdout: bool, - conn: 'Connection', + name: str, ) -> Process: """Create and start a Hathor node process.""" - run_hathor_node_args = (argv, RunNode, logging_output, capture_stdout, conn) - hathor_node_process = Process(target=_run_hathor_node, args=run_hathor_node_args) - hathor_node_process.start() - return hathor_node_process + args = _RunNodeArgs( + argv=argv, + runner=runner, + logging_output=logging_output, + capture_stdout=capture_stdout, + name=name + ) + process = Process(target=_run_node, args=(args,), name=name) + process.start() + return process -def _run_hathor_node( - argv: list[str], - run_node_cmd: type[RunNode], - logging_output: LoggingOutput, - capture_stdout: bool, - conn: 'Connection', -) -> None: - """Function to be called by a background process to run the Hathor full node.""" - from hathor.cli.util import process_logging_options, setup_logging +@dataclass(frozen=True, slots=True, kw_only=True) +class _RunNodeArgs: + argv: list[str] + runner: type[RunNode] + logging_output: LoggingOutput + capture_stdout: bool + name: str - # We don't terminate via SIGINT directly, instead the side-dag process will terminate us. - signal.signal(signal.SIGINT, lambda _, __: None) + +def _run_node(args: _RunNodeArgs) -> None: + from hathor.cli.util import process_logging_options, setup_logging try: - log_options = process_logging_options(argv) + log_options = process_logging_options(args.argv) setup_logging( - logging_output=logging_output, + logging_output=args.logging_output, logging_options=log_options, - capture_stdout=capture_stdout, - extra_log_info=_get_extra_log_info('hathor') + capture_stdout=args.capture_stdout, + extra_log_info=_get_extra_log_info(args.name) ) - hathor_node = run_node_cmd(argv=argv) - except (BaseException, Exception): - logger.exception('hathor process terminated due to exception in initialization') - conn.send(HathorProcessInitFail(traceback.format_exc())) + logger.info(f'initializing node "{args.name}"') + node = args.runner(argv=args.argv) + except KeyboardInterrupt: + logger.warn(f'{args.name} node interrupted by user') return - - conn.send(HathorProcessInitSuccess()) - hathor_node.run() - - # If `run()` returns, either the side-dag node exited and terminated us, leaving the message below, or the hathor - # node exited and we will terminate the side-dag node. - if conn.poll(): - message = conn.recv() - assert isinstance(message, SideDagProcessTerminated) - logger.critical('hathor node terminated because side-dag process was terminated') + except (BaseException, Exception): + logger.exception(f'process "{args.name}" terminated due to exception in initialization') return - conn.send(HathorProcessTerminated()) - logger.critical('terminating side-dag node...') - parent_pid = os.getppid() - parent_process = psutil.Process(parent_pid) - parent_process.terminate() + node.run() + logger.critical(f'node "{args.name}" gracefully terminated') def _get_extra_log_info(process_name: str) -> dict[str, str]: diff --git a/tests/cli/test_side_dag.py b/tests/cli/test_side_dag.py index 5f799f1275..0d4caf4a5c 100644 --- a/tests/cli/test_side_dag.py +++ b/tests/cli/test_side_dag.py @@ -12,22 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import ANY, Mock, call, patch - import pytest -from hathor.cli import side_dag -from hathor.cli.side_dag import ( - HathorProcessInitFail, - HathorProcessInitSuccess, - HathorProcessTerminated, - SideDagProcessTerminated, - _partition_argv, - _run_hathor_node, - _run_side_dag_node, - main, -) -from hathor.cli.util import LoggingOutput +from hathor.cli.side_dag import _partition_argv @pytest.mark.parametrize( @@ -50,7 +37,7 @@ ), ] ) -def test_process_argv( +def test_partition_argv( argv: list[str], expected_hathor_node_argv: list[str], expected_side_dag_argv: list[str] @@ -59,127 +46,3 @@ def test_process_argv( assert hathor_node_argv == expected_hathor_node_argv assert side_dag_argv == expected_side_dag_argv - - -def test_run_side_dag_node_hathor_init_failed() -> None: - argv: list[str] = [] - conn_mock = Mock() - conn_mock.poll = Mock(return_value=True) - conn_mock.recv = Mock(return_value=HathorProcessInitFail('some reason')) - hathor_node_process = Mock() - - with patch.object(side_dag, 'SideDagRunNode') as side_dag_mock: - _run_side_dag_node(argv, conn=conn_mock, hathor_node_process=hathor_node_process) - side_dag_mock.assert_not_called() - hathor_node_process.terminate.assert_not_called() - conn_mock.send.assert_not_called() - - -def test_run_side_dag_node_init_failed() -> None: - argv: list[str] = [] - conn_mock = Mock() - conn_mock.poll = Mock(return_value=True) - conn_mock.recv = Mock(return_value=HathorProcessInitSuccess()) - hathor_node_process = Mock() - side_dag_mock = Mock(side_effect=Exception) - - with patch.object(side_dag, 'SideDagRunNode', side_dag_mock): - _run_side_dag_node(argv, conn=conn_mock, hathor_node_process=hathor_node_process) - side_dag_mock.assert_called_once_with(argv=argv) - hathor_node_process.terminate.assert_called_once() - conn_mock.send.assert_called_once_with(SideDagProcessTerminated()) - - -def test_run_side_dag_node_hathor_terminated() -> None: - argv: list[str] = [] - conn_mock = Mock() - conn_mock.poll = Mock(side_effect=[True, True]) - conn_mock.recv = Mock(side_effect=[HathorProcessInitSuccess(), HathorProcessTerminated()]) - hathor_node_process = Mock() - side_dag_instance_mock = Mock() - side_dag_mock = Mock(return_value=side_dag_instance_mock) - - with patch.object(side_dag, 'SideDagRunNode', side_dag_mock): - _run_side_dag_node(argv, conn=conn_mock, hathor_node_process=hathor_node_process) - side_dag_mock.assert_called_once_with(argv=argv) - side_dag_instance_mock.run.assert_called_once() - hathor_node_process.terminate.assert_not_called() - conn_mock.send.assert_not_called() - - -def test_run_side_dag_node_terminated() -> None: - argv: list[str] = [] - conn_mock = Mock() - conn_mock.poll = Mock(side_effect=[True, False]) - conn_mock.recv = Mock(side_effect=[HathorProcessInitSuccess()]) - hathor_node_process = Mock() - side_dag_instance_mock = Mock() - side_dag_mock = Mock(return_value=side_dag_instance_mock) - - with patch.object(side_dag, 'SideDagRunNode', side_dag_mock): - _run_side_dag_node(argv, conn=conn_mock, hathor_node_process=hathor_node_process) - side_dag_mock.assert_called_once_with(argv=argv) - side_dag_instance_mock.run.assert_called_once() - hathor_node_process.terminate.assert_called_once() - conn_mock.send.assert_called_once_with(SideDagProcessTerminated()) - - -def test_run_hathor_node_init_failed() -> None: - argv: list[str] = [] - run_node_cmd_mock = Mock(side_effect=Exception) - capture_stdout = False - conn_mock = Mock() - parent_process_mock = Mock() - - with patch('psutil.Process', return_value=parent_process_mock): - _run_hathor_node(argv, run_node_cmd_mock, LoggingOutput.PRETTY, capture_stdout, conn_mock) - run_node_cmd_mock.assert_called_once_with(argv=argv) - conn_mock.send.assert_called_once_with(HathorProcessInitFail(ANY)) - parent_process_mock.terminate.assert_not_called() - - -def test_run_hathor_node_side_dag_terminated() -> None: - argv: list[str] = [] - run_node_instance = Mock() - run_node_cmd_mock = Mock(return_value=run_node_instance) - capture_stdout = False - conn_mock = Mock() - conn_mock.poll = Mock(return_value=True) - conn_mock.recv = Mock(return_value=SideDagProcessTerminated()) - parent_process_mock = Mock() - - with patch('psutil.Process', return_value=parent_process_mock): - _run_hathor_node(argv, run_node_cmd_mock, LoggingOutput.PRETTY, capture_stdout, conn_mock) - conn_mock.send.assert_called_once_with(HathorProcessInitSuccess()) - run_node_instance.run.assert_called_once() - parent_process_mock.terminate.assert_not_called() - - -def test_run_hathor_node_exited() -> None: - argv: list[str] = [] - run_node_instance = Mock() - run_node_cmd_mock = Mock(return_value=run_node_instance) - capture_stdout = False - conn_mock = Mock() - conn_mock.poll = Mock(return_value=False) - parent_process_mock = Mock() - - with patch('psutil.Process', return_value=parent_process_mock): - _run_hathor_node(argv, run_node_cmd_mock, LoggingOutput.PRETTY, capture_stdout, conn_mock) - conn_mock.send.assert_has_calls([call(HathorProcessInitSuccess()), call(HathorProcessTerminated())]) - run_node_instance.run.assert_called_once() - parent_process_mock.terminate.assert_called_once() - - -def dummy_run_hathor_node() -> None: - pass - - -def test_main_success() -> None: - with ( - patch.object(side_dag, '_run_side_dag_node') as run_side_dag_mock, - patch.object(side_dag, '_run_hathor_node', dummy_run_hathor_node), - - ): - main(capture_stdout=False) - run_side_dag_mock.assert_called_once()