diff --git a/RELEASE.md b/RELEASE.md index 5be798e33..e7e50b88d 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -22,6 +22,7 @@ Please follow the established format: - Enable SQLite WAL mode for Azure ML to fix database locking issues (#2131) - Replace `flake8`, `isort`, `pylint` and `black` by `ruff` (#2149) - Refactor `DatasetStatsHook` to avoid showing error when dataset doesn't have file size info (#2174) +- Add check for port availability before starting Kedro Viz to prevent unintended browser redirects when the port is already in use (#2176) # Release 10.0.0 diff --git a/package/kedro_viz/launchers/cli/run.py b/package/kedro_viz/launchers/cli/run.py index b2e74a48b..e4093b940 100644 --- a/package/kedro_viz/launchers/cli/run.py +++ b/package/kedro_viz/launchers/cli/run.py @@ -115,6 +115,7 @@ def run( from kedro_viz.launchers.utils import ( _PYPROJECT, _check_viz_up, + _find_available_port, _find_kedro_project, _start_browser, _wait_for, @@ -145,6 +146,9 @@ def run( "https://github.com/kedro-org/kedro-viz/releases.", "yellow", ) + + port = _find_available_port(host, port) + try: if port in _VIZ_PROCESSES and _VIZ_PROCESSES[port].is_alive(): _VIZ_PROCESSES[port].terminate() @@ -186,7 +190,6 @@ def run( ) display_cli_message("Starting Kedro Viz ...", "green") - viz_process.start() _VIZ_PROCESSES[port] = viz_process diff --git a/package/kedro_viz/launchers/utils.py b/package/kedro_viz/launchers/utils.py index 5c6bbae9e..50f8e6e84 100644 --- a/package/kedro_viz/launchers/utils.py +++ b/package/kedro_viz/launchers/utils.py @@ -2,6 +2,8 @@ used in the `kedro_viz.launchers` package.""" import logging +import socket +import sys import webbrowser from pathlib import Path from time import sleep, time @@ -80,6 +82,33 @@ def _check_viz_up(host: str, port: int): return response.status_code == 200 +def _is_port_in_use(host: str, port: int): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + return s.connect_ex((host, port)) == 0 + + +def _find_available_port(host: str, start_port: int, max_attempts: int = 5) -> int: + max_port = start_port + max_attempts - 1 + port = start_port + while port <= max_port: + if not _is_port_in_use(host, port): + return port + display_cli_message( + f"Port {port} is already in use. Trying the next port...", + "yellow", + ) + port += 1 + display_cli_message( + f"Error: All ports in the range {start_port}-{max_port} are in use.", + "red", + ) + display_cli_message( + "Please specify a different port using the '--port' option.", + "red", + ) + sys.exit(1) + + def _is_localhost(host: str) -> bool: """Check whether a host is a localhost""" return host in ("127.0.0.1", "localhost", "0.0.0.0") diff --git a/package/tests/test_launchers/test_cli/test_run.py b/package/tests/test_launchers/test_cli/test_run.py index 86adae92f..95a809d2e 100644 --- a/package/tests/test_launchers/test_cli/test_run.py +++ b/package/tests/test_launchers/test_cli/test_run.py @@ -10,7 +10,7 @@ from kedro_viz.autoreload_file_filter import AutoreloadFileFilter from kedro_viz.launchers.cli import main from kedro_viz.launchers.cli.run import _VIZ_PROCESSES -from kedro_viz.launchers.utils import _PYPROJECT +from kedro_viz.launchers.utils import _PYPROJECT, _find_available_port from kedro_viz.server import run_server @@ -217,6 +217,9 @@ def test_kedro_viz_command_run_server( "kedro_viz.launchers.utils._wait_for.__defaults__", (True, 1, True, 1) ) + # Mock _is_port_in_use to speed up test. + mocker.patch("kedro_viz.launchers.utils._is_port_in_use", return_value=False) + # Mock finding kedro project mocker.patch( "kedro_viz.launchers.utils._find_kedro_project", @@ -394,3 +397,17 @@ def test_kedro_viz_command_with_autoreload( kwargs={**run_process_kwargs}, ) assert run_process_kwargs["kwargs"]["port"] in _VIZ_PROCESSES + + # Test case to simulate port occupation and check available port selection + def test_find_available_port_with_occupied_ports(self, mocker): + mock_is_port_in_use = mocker.patch("kedro_viz.launchers.utils._is_port_in_use") + + # Mock ports 4141, 4142 being occupied and 4143 is free + mock_is_port_in_use.side_effect = [True, True, False] + + available_port = _find_available_port("127.0.0.1", 4141) + + # Assert that the function returns the first free port, 4143 + assert ( + available_port == 4143 + ), "Expected port 4143 to be returned as the available port" diff --git a/package/tests/test_launchers/test_utils.py b/package/tests/test_launchers/test_utils.py index 83e9203bd..fd2043af7 100644 --- a/package/tests/test_launchers/test_utils.py +++ b/package/tests/test_launchers/test_utils.py @@ -7,6 +7,7 @@ from kedro_viz.launchers.utils import ( _check_viz_up, + _find_available_port, _find_kedro_project, _is_project, _start_browser, @@ -99,3 +100,20 @@ def test_toml_bad_encoding(self, mocker): def test_find_kedro_project(project_dir, is_project_found, expected, mocker): mocker.patch("kedro_viz.launchers.utils._is_project", return_value=is_project_found) assert _find_kedro_project(Path(project_dir)) == expected + + +def test_find_available_port_all_ports_occupied(mocker): + mocker.patch("kedro_viz.launchers.utils._is_port_in_use", return_value=True) + mock_display_message = mocker.patch("kedro_viz.launchers.utils.display_cli_message") + + # Check for SystemExit when all ports are occupied + with pytest.raises(SystemExit) as exit_exception: + _find_available_port("127.0.0.1", 4141, max_attempts=5) + assert exit_exception.value.code == 1 + + mock_display_message.assert_any_call( + "Error: All ports in the range 4141-4145 are in use.", "red" + ) + mock_display_message.assert_any_call( + "Please specify a different port using the '--port' option.", "red" + )