diff --git a/src/isolate/connections/grpc/_base.py b/src/isolate/connections/grpc/_base.py index b258de0..7ede914 100644 --- a/src/isolate/connections/grpc/_base.py +++ b/src/isolate/connections/grpc/_base.py @@ -145,14 +145,28 @@ def find_free_port() -> Tuple[str, int]: def abort_agent(self) -> None: if self._process is not None: + return_code: int | None = None try: - print("Terminating the agent process...") - self._process.terminate() - self._process.wait(timeout=PROCESS_SHUTDOWN_TIMEOUT_SECONDS) - print("Agent process shutdown gracefully") + if self._process.poll() is not None: + # already finished + return_code = self._process.returncode + else: + print("Terminating the agent process...") + self._process.terminate() + return_code = self._process.wait( + timeout=PROCESS_SHUTDOWN_TIMEOUT_SECONDS + ) + print("Agent process shutdown gracefully") except Exception as exc: print(f"Failed to shutdown the agent process gracefully: {exc}") self._process.kill() + return_code = self._process.wait() + + self.log( + f"Isolate agent finished (exit code: {return_code})", + level=LogLevel.INFO, + source=LogSource.BRIDGE, + ) self._process = None def is_alive(self) -> bool: diff --git a/tests/test_connections_grpc_base.py b/tests/test_connections_grpc_base.py new file mode 100644 index 0000000..7e6e249 --- /dev/null +++ b/tests/test_connections_grpc_base.py @@ -0,0 +1,77 @@ +from pathlib import Path +from unittest.mock import Mock + +from isolate.backends.local import LocalPythonEnvironment +from isolate.connections import LocalPythonGRPC +from isolate.logs import LogLevel, LogSource + + +def make_connection(tmp_path: Path) -> LocalPythonGRPC: + environment = LocalPythonEnvironment() + return LocalPythonGRPC(environment, tmp_path) + + +def test_abort_agent_logs_return_code_for_already_exited_process( + tmp_path: Path, +) -> None: + connection = make_connection(tmp_path) + process = Mock() + process.poll.return_value = 0 + process.returncode = 0 + connection._process = process + + connection.log = Mock() + connection.abort_agent() + + process.terminate.assert_not_called() + process.wait.assert_not_called() + process.kill.assert_not_called() + connection.log.assert_called_once_with( + "Isolate agent finished (exit code: 0)", + level=LogLevel.INFO, + source=LogSource.BRIDGE, + ) + assert connection._process is None + + +def test_abort_agent_logs_return_code_for_graceful_termination(tmp_path: Path) -> None: + connection = make_connection(tmp_path) + process = Mock() + process.poll.return_value = None + process.wait.return_value = -15 + connection._process = process + + connection.log = Mock() + connection.abort_agent() + + process.terminate.assert_called_once() + process.wait.assert_called_once() + process.kill.assert_not_called() + connection.log.assert_called_once_with( + "Isolate agent finished (exit code: -15)", + level=LogLevel.INFO, + source=LogSource.BRIDGE, + ) + assert connection._process is None + + +def test_abort_agent_logs_return_code_after_kill_fallback(tmp_path: Path) -> None: + connection = make_connection(tmp_path) + process = Mock() + process.poll.return_value = None + process.terminate.side_effect = RuntimeError("terminate failed") + process.wait.return_value = -9 + connection._process = process + + connection.log = Mock() + connection.abort_agent() + + process.terminate.assert_called_once() + process.kill.assert_called_once() + process.wait.assert_called_once() + connection.log.assert_called_once_with( + "Isolate agent finished (exit code: -9)", + level=LogLevel.INFO, + source=LogSource.BRIDGE, + ) + assert connection._process is None