From 4b012ffd8c2f2f34e0ee31476fe1fc3681a0cc7a Mon Sep 17 00:00:00 2001 From: cotran Date: Thu, 24 Oct 2024 15:29:22 -0700 Subject: [PATCH] reformat --- model_server/app/cli.py | 8 ++++++-- model_server/app/tests/test_cli_stop_server.py | 8 ++++---- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/model_server/app/cli.py b/model_server/app/cli.py index 425dd23e..dd6a5679 100644 --- a/model_server/app/cli.py +++ b/model_server/app/cli.py @@ -98,12 +98,11 @@ def check_and_install_lsof(): print(f"Failed to install lsof: {install_error}") -def stop_server(port=51000, wait=True, timeout=10): +def kill_process(port=51000, wait=True, timeout=10): """Stop the running Uvicorn server.""" log.info("Stopping model server") try: # Run the function to check and install lsof if necessary - check_and_install_lsof() # Step 1: Run lsof command to get the process using the port lsof_command = f"lsof -n | grep {port} | grep -i LISTEN" result = subprocess.run( @@ -156,6 +155,11 @@ def stop_server(port=51000, wait=True, timeout=10): print(f"Error occurred: {e}") +def stop_server(port=51000, wait=True, timeout=10): + check_and_install_lsof() + kill_process(port, wait, timeout) + + def restart_server(port=51000): """Restart the Uvicorn server.""" stop_server(port) diff --git a/model_server/app/tests/test_cli_stop_server.py b/model_server/app/tests/test_cli_stop_server.py index 5b80303b..4f3955a7 100644 --- a/model_server/app/tests/test_cli_stop_server.py +++ b/model_server/app/tests/test_cli_stop_server.py @@ -2,7 +2,7 @@ import unittest from unittest.mock import patch, MagicMock import subprocess import time -from app.cli import stop_server +from app.cli import kill_process class TestStopServer(unittest.TestCase): @@ -11,7 +11,7 @@ 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: - stop_server(port=51000) + kill_process(port=51000) mock_print.assert_called_with("No process found listening on port 51000.") @patch("subprocess.run") @@ -23,7 +23,7 @@ class TestStopServer(unittest.TestCase): MagicMock(returncode=1), # for checking the process after it is killed ] with patch("builtins.print") as mock_print: - stop_server(port=51000, wait=True, timeout=5) + 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.") @@ -42,7 +42,7 @@ class TestStopServer(unittest.TestCase): ] with patch("builtins.print") as mock_print: - stop_server(port=51000, wait=True, timeout=5) + kill_process(port=51000, 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")