Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions src/isolate/connections/grpc/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
77 changes: 77 additions & 0 deletions tests/test_connections_grpc_base.py
Original file line number Diff line number Diff line change
@@ -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