Add port in use check to prevent browser redirecting incorrectly for …

…kedro viz (#2176)

* add check if port is in use

Signed-off-by: Sajid Alam <sajid_alam@mckinsey.com>

* Update RELEASE.md

Signed-off-by: Sajid Alam <sajid_alam@mckinsey.com>

* automatically increment port

Signed-off-by: Sajid Alam <sajid_alam@mckinsey.com>

* fix unit tests

Signed-off-by: Sajid Alam <sajid_alam@mckinsey.com>

* add port occupied test

Signed-off-by: Sajid Alam <sajid_alam@mckinsey.com>

* coverage

Signed-off-by: Sajid Alam <sajid_alam@mckinsey.com>

* speed up test by mocking _is_port_in_use to return false

Signed-off-by: Sajid Alam <sajid_alam@mckinsey.com>

---------

Signed-off-by: Sajid Alam <sajid_alam@mckinsey.com>
Signed-off-by: Sajid Alam <90610031+SajidAlamQB@users.noreply.github.com>

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

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

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")

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"

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"
+    )