This commit is contained in:
cotran 2024-12-10 19:31:02 -08:00
parent 2405fb36e3
commit b2ef3f7266
4 changed files with 32 additions and 39 deletions

View file

@ -1,7 +1,6 @@
import unittest
from unittest.mock import patch, MagicMock
from src.core.cli import kill_process
from src.commons.utils import kill_processes
class TestStopServer(unittest.TestCase):
@ -10,8 +9,8 @@ class TestStopServer(unittest.TestCase):
# Mock subprocess.run to simulate no process listening on the port
mock_run.return_value.returncode = 1
with patch("builtins.print") as mock_print:
kill_process(port=51000)
mock_print.assert_called_with("No process found listening on port 51000.")
kill_processes(port_processes=[""], wait=True, timeout=5)
mock_print.assert_not_called()
@patch("subprocess.run")
def test_stop_server_process_killed(self, mock_run):
@ -22,18 +21,15 @@ class TestStopServer(unittest.TestCase):
MagicMock(returncode=1), # for checking the process after it is killed
]
with patch("builtins.print") as mock_print:
kill_process(port=51000, wait=True, timeout=5)
mock_print.assert_any_call("Killing model server process with PID 1234")
mock_print.assert_any_call("Process 1234 has been killed.")
kill_processes(
port_processes=["uvicorn 1234 user LISTEN\n"], wait=True, timeout=5
)
mock_print.assert_any_call("Killing process with PID 1234...")
@patch("subprocess.run")
def test_stop_server_multiple_pids(self, mock_run):
# Simulate lsof returning multiple process ids (e.g., 1234 and 5678)
mock_run.side_effect = [
MagicMock(
returncode=0,
stdout="uvicorn 1234 user LISTEN\nuvicorn 5678 user LISTEN\n",
), # lsof output
MagicMock(returncode=0), # first kill command for PID 1234
MagicMock(returncode=1), # PID 1234 is successfully terminated
MagicMock(returncode=0), # second kill command for PID 5678
@ -41,13 +37,15 @@ class TestStopServer(unittest.TestCase):
]
with patch("builtins.print") as mock_print:
kill_process(port=51000, wait=True, timeout=5)
kill_processes(
port_processes=["uvicorn 1234 user LISTEN", "uvicorn 5678 user LISTEN"],
wait=True,
timeout=5,
)
# Assert that the function tried to kill both PIDs
mock_print.assert_any_call("Killing model server process with PID 1234")
mock_print.assert_any_call("Process 1234 has been killed.")
mock_print.assert_any_call("Killing model server process with PID 5678")
mock_print.assert_any_call("Process 5678 has been killed.")
mock_print.assert_any_call("Killing process with PID 1234...")
mock_print.assert_any_call("Killing process with PID 5678...")
if __name__ == "__main__":