Block launch rocky server multiple times

src/ansys/rocky/core/client.py

@@ -21,18 +21,23 @@
 # SOFTWARE.
 
 import pickle
-from typing import Generator
+from typing import Final, Generator
 
 import Pyro5.api
 import serpent
 
 from ansys.rocky.core.exceptions import RockyApiError
 
+DEFAULT_SERVER_PORT: Final[int] = 50615
+
 _ROCKY_API = None
 
 
-def connect_to_rocky(host: str = "localhost", port: int = 50615) -> "RockyClient":
-    """Connect to a Rocky Application instance.
+def connect_to_rocky(
+    host: str = "localhost", port: int = DEFAULT_SERVER_PORT
+) -> "RockyClient":
+    """
+    Connect to a Rocky Application instance.
 
     Parameters
     ----------

src/ansys/rocky/core/launcher.py

@@ -25,17 +25,21 @@
 from pathlib import Path
 import subprocess
 import time
-from typing import Optional
+from typing import Optional, Union
 
 from Pyro5.errors import CommunicationError
 
-from ansys.rocky.core.client import RockyClient, connect_to_rocky
+from ansys.rocky.core.client import DEFAULT_SERVER_PORT, RockyClient, connect_to_rocky
 from ansys.rocky.core.exceptions import RockyLaunchError
 
+_CONNECT_TO_SERVER_TIMEOUT = 60
+
 
 def launch_rocky(
-    rocky_exe: Optional[Path] = None,
+    rocky_exe: Optional[Union[Path, str]] = None,
+    *,
     headless: bool = True,
+    server_port: int = DEFAULT_SERVER_PORT,
 ) -> RockyClient:
     """
     Launch Rocky executable with PyRocky server enabled, wait Rocky to start up and
@@ -48,12 +52,20 @@ def launch_rocky(
         environment variables `AWP_ROOT241` and `AWP_ROOT232`.
     headless : bool, optional
         Whether to launch Rocky in headless mode. Default is `True`.
+    server_port: int, optional
+        Set the port used to host Rocky PyRocky server.
 
     Returns
     -------
     RockyClient
         A `RockyClient` instance connected to the launched Rocky application.
     """
+    if isinstance(rocky_exe, str):
+        rocky_exe = Path(rocky_exe)
+
+    if _is_port_busy(server_port):
+        raise RockyLaunchError(f"Port {server_port} already in use")
+
     if rocky_exe is None:
         for awp_root in ["AWP_ROOT241", "AWP_ROOT232"]:
             if awp_root not in os.environ:
@@ -68,7 +80,7 @@ def launch_rocky(
         if not rocky_exe.is_file():
             raise FileNotFoundError(f"Rocky executable not found at {rocky_exe}")
 
-    cmd = [rocky_exe, "--pyrocky"]
+    cmd = [rocky_exe, "--pyrocky", "--pyrocky-port", str(server_port)]
     if headless:
         cmd.append("--headless")
     with contextlib.suppress(subprocess.TimeoutExpired):
@@ -79,10 +91,11 @@ def launch_rocky(
     if rocky_process.returncode is not None:
         raise RockyLaunchError(f"Error launching Rocky:\n  {' '.join(cmd)}")
 
-    client = connect_to_rocky()
+    client = connect_to_rocky(port=server_port)
 
     # TODO: A more elegant way to find out that Rocky Pyro server started.
-    for _ in range(_WAIT_ROCKY_START):
+    now = time.time()
+    while (time.time() - now) < _CONNECT_TO_SERVER_TIMEOUT:
         try:
             client.api.GetProject()
         except CommunicationError:
@@ -96,4 +109,21 @@ def launch_rocky(
     return client
 
 
-_WAIT_ROCKY_START = 60
+def _is_port_busy(port: int) -> bool:
+    """
+    Check if there is already a Rocky server running.
+
+    Parameters
+    ----------
+    port : int
+        The port to be checked.
+
+    Returns
+    -------
+    bool
+        Whether the port is busy or not.
+    """
+    import socket
+
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        return s.connect_ex(("localhost", port)) == 0

tests/test_pyrocky.py

@@ -23,6 +23,8 @@
 import pytest
 
 import ansys.rocky.core as pyrocky
+from ansys.rocky.core.client import DEFAULT_SERVER_PORT
+from ansys.rocky.core.launcher import RockyLaunchError
 
 
 @pytest.fixture()
@@ -81,3 +83,18 @@ def test_sequences_interface(rocky_session):
     # Test __del__
     del inlets_outlets[0]
     assert {e.GetName() for e in inlets_outlets} == {"Inlet2"}
+
+
+def test_pyrocky_launch_multiple_servers():
+    """
+    Test that start multiple rocky servers is not allowed.
+    """
+    import socket
+
+    # Emulating Rocky server already running by binding socket to the server address.
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("localhost", DEFAULT_SERVER_PORT))
+        s.listen(10)
+
+        with pytest.raises(RockyLaunchError, match=r"Port \d+ already in use"):
+            pyrocky.launch_rocky()