Enhance RemoteClient connection logic with exponential backoff and update method documentation; add comprehensive unit tests for RemoteProxy and SimpleRemoteProxy interactions

PierreRaybaut · PierreRaybaut · commit a6323909a291 · 2025-10-06T10:39:52.000+02:00
diff --git a/datalab/control/remote.py b/datalab/control/remote.py
@@ -753,32 +753,50 @@ def connect(
         Args:
             port: XML-RPC port to connect to. If not specified,
              the port is automatically retrieved from DataLab configuration.
-            timeout: Timeout in seconds. Defaults to 5.0.
-            retries: Number of retries. Defaults to 10.
+            timeout: Maximum time to wait for connection in seconds. Defaults to 5.0.
+             This is the total maximum wait time, not per retry.
+            retries: Number of retries. Defaults to 10. This parameter is deprecated
+             and will be removed in a future version (kept for backward compatibility).
 
         Raises:
             ConnectionRefusedError: Unable to connect to DataLab
             ValueError: Invalid timeout (must be >= 0.0)
             ValueError: Invalid number of retries (must be >= 1)
         """
         timeout = 5.0 if timeout is None else timeout
-        retries = 10 if retries is None else retries
+        retries = 10 if retries is None else retries  # Kept for backward compatibility
         if timeout < 0.0:
             raise ValueError("timeout must be >= 0.0")
         if retries < 1:
             raise ValueError("retries must be >= 1")
-        self.set_port(port)
+
         execenv.print(f"Connecting to DataLab XML-RPC server... [port:{port}] ", end="")
-        for _index in range(retries):
+
+        # Use exponential backoff for more efficient retrying
+        start_time = time.time()
+        poll_interval = 0.1  # Start with 100ms
+        max_poll_interval = 1.0  # Cap at 1 second
+
+        while True:
             try:
+                # Try to set the port - this may fail if DataLab hasn't written its
+                # config file yet, so we retry it in the loop
+                self.set_port(port)
                 self.__connect_to_server()
-                break
-            except ConnectionRefusedError:
-                time.sleep(timeout / retries)
-        else:
-            execenv.print("KO")
-            raise ConnectionRefusedError("Unable to connect to DataLab")
-        execenv.print(f"OK (port: {self.port})")
+                elapsed = time.time() - start_time
+                execenv.print(f"OK (port: {self.port}, connected in {elapsed:.1f}s)")
+                return
+            except (ConnectionRefusedError, OSError):
+                # Catch both ConnectionRefusedError and OSError (includes socket errors)
+                elapsed = time.time() - start_time
+                if elapsed >= timeout:
+                    execenv.print("KO")
+                    raise ConnectionRefusedError(
+                        f"Unable to connect to DataLab after {elapsed:.1f}s"
+                    )
+                # Wait before next retry with exponential backoff
+                time.sleep(poll_interval)
+                poll_interval = min(poll_interval * 1.5, max_poll_interval)
 
     def disconnect(self) -> None:
         """Disconnect from DataLab XML-RPC server."""
diff --git a/datalab/tests/__init__.py b/datalab/tests/__init__.py
@@ -34,6 +34,7 @@
 
 import datalab.config  # Loading icons
 from datalab.config import MOD_NAME
+from datalab.control.proxy import RemoteProxy
 from datalab.env import execenv
 from datalab.gui.main import DLMainWindow
 from datalab.gui.panel.image import ImagePanel
@@ -107,7 +108,7 @@ def is_pid_alive(pid: int) -> bool:
     return psutil.pid_exists(pid) and psutil.Process(pid).is_running()
 
 
-def run_datalab_in_background():
+def run_datalab_in_background(wait_until_ready: bool = True) -> None:
     """Run DataLab application as a service.
 
     This function starts the DataLab application in a separate process, ensuring that
@@ -121,8 +122,13 @@ def run_datalab_in_background():
     application needs to be running in the background while a client connects to it
     and performs various operations.
 
+    Args:
+        wait_until_ready: If True, waits until the DataLab application is ready to
+         accept connections (default: True). Uses RemoteProxy's built-in retry logic
+         with extended timeout to handle DataLab startup time.
+
     Raises:
-        AssertionError: If the DataLab application fails to start
+        RuntimeError: If the DataLab application fails to start
     """
     env = os.environ.copy()
     env[execenv.DO_NOT_QUIT_ENV] = "1"
@@ -141,8 +147,26 @@ def run_datalab_in_background():
     # error messages or logs that might help you understand why the application failed
     # to start.
     # If the script is executed within a pytest session, add the `-s` option to pytest.
-    time.sleep(2)
-    assert is_pid_alive(proc.pid), "Unable to start DataLab application"
+
+    # Give the process a moment to actually start
+    time.sleep(1)
+    if not is_pid_alive(proc.pid):
+        raise RuntimeError("DataLab process terminated immediately after start")
+
+    if wait_until_ready:
+        # Use RemoteProxy's built-in retry mechanism with extended timeout
+        # DataLab startup: Python imports, Qt init, GUI creation, XML-RPC server
+        try:
+            proxy = RemoteProxy(autoconnect=False)
+            proxy.connect(timeout=30.0)  # 30 seconds max for DataLab to be ready
+            proxy.disconnect()
+        except ConnectionRefusedError as exc:
+            if is_pid_alive(proc.pid):
+                proc.kill()
+            raise RuntimeError(
+                "Failed to connect to DataLab application. "
+                "Process may have started but XML-RPC server is not responding."
+            ) from exc
 
 
 @contextmanager
diff --git a/datalab/tests/features/control/simpleclient_unit_test.py b/datalab/tests/features/control/simpleclient_unit_test.py
@@ -0,0 +1,184 @@
+# Copyright (c) DataLab Platform Developers, BSD 3-Clause license, see LICENSE file.
+
+"""
+Simple Remote client test
+-------------------------
+
+This test covers the simple remote client functionalities as provided by the
+`sigima.client` module.
+
+In Sigima, the `sigima.tests.common.client_unit_test` module provides a set of unit
+tests for the client functionalities.
+
+The purpose of these tests is to ensure that the client can correctly interact with
+the real server, handling various scenarios and edge cases.
+
+The tests include:
+
+- A function comparing the list of methods implemented by the real server
+  (DataLab's `RemoteServer` class) to those implemented by the stub server
+  (Sigima's `DataLabStubServer` class).
+- A function comparing the list of methods implemented by the full client (DataLab's
+  `RemoteProxy` class) to those implemented by the simple client (Sigima's
+  `SimpleRemoteProxy` class).
+- A function that simply runs the `sigima.tests.common.client_unit_test` suite after
+  having launched a real server instance (DataLab application, using the
+  `datalab.tests.run_datalab_in_background` function), ensuring that this test passes
+  with the stub server as well as with the real server.
+"""
+
+# pylint: disable=invalid-name  # Allows short reference names like x, y, ...
+# guitest: skip
+
+import inspect
+
+from guidata.env import execenv
+from sigima.client.remote import SimpleRemoteProxy
+from sigima.client.stub import DataLabStubServer
+from sigima.tests.common.client_unit_test import RemoteClientTester
+
+from datalab.control.baseproxy import AbstractDLControl
+from datalab.control.proxy import RemoteProxy
+from datalab.tests import run_datalab_in_background
+
+
+def test_compare_server_methods() -> None:
+    """Compare methods implemented by RemoteServer vs DataLabStubServer.
+
+    This function compares the list of methods implemented by the real server
+    (DataLab's RemoteServer class) to those implemented by the stub server
+    (Sigima's DataLabStubServer class).
+    """
+    execenv.print("\n=== Comparing Server Methods ===")
+
+    # Get methods from RemoteServer (the real DataLab server)
+    # RemoteServer implements AbstractDLControl methods
+    remote_server_methods = set(AbstractDLControl.get_public_methods())
+
+    # Get methods from DataLabStubServer (the stub server)
+    stub_server_methods = {
+        name
+        for name, method in inspect.getmembers(DataLabStubServer, inspect.isfunction)
+        if not name.startswith("_")
+    }
+
+    execenv.print(f"RemoteServer methods: {len(remote_server_methods)}")
+    execenv.print(f"DataLabStubServer methods: {len(stub_server_methods)}")
+
+    # Methods in RemoteServer but not in DataLabStubServer
+    missing_in_stub = remote_server_methods - stub_server_methods
+    if missing_in_stub:
+        execenv.print("\n⚠️  Methods in RemoteServer but missing in DataLabStubServer:")
+        for method in sorted(missing_in_stub):
+            execenv.print(f"  - {method}")
+
+    # Methods in DataLabStubServer but not in RemoteServer
+    extra_in_stub = stub_server_methods - remote_server_methods
+    if extra_in_stub:
+        execenv.print("\n🔍 Methods in DataLabStubServer but not in RemoteServer:")
+        for method in sorted(extra_in_stub):
+            execenv.print(f"  - {method}")
+
+    # Common methods
+    common_methods = remote_server_methods & stub_server_methods
+    execenv.print(f"\n✅ Common methods: {len(common_methods)}")
+
+    # Assert that stub server implements all required methods
+    assert not missing_in_stub, (
+        f"DataLabStubServer is missing {len(missing_in_stub)} methods "
+        f"implemented by RemoteServer: {sorted(missing_in_stub)}"
+    )
+
+    execenv.print("✨ Server method comparison completed successfully!")
+
+
+def test_compare_client_methods() -> None:
+    """Compare methods implemented by RemoteProxy vs SimpleRemoteProxy.
+
+    This function compares the list of methods implemented by the full client
+    (DataLab's RemoteProxy class) to those implemented by the simple client
+    (Sigima's SimpleRemoteProxy class).
+    """
+    execenv.print("\n=== Comparing Client Methods ===")
+
+    # Get public methods from RemoteProxy (the full client)
+    remote_proxy_methods = {
+        name
+        for name, method in inspect.getmembers(RemoteProxy, inspect.ismethod)
+        if not name.startswith("_") and callable(method)
+    }
+
+    # Get public methods from SimpleRemoteProxy (the simple client)
+    simple_proxy_methods = {
+        name
+        for name, method in inspect.getmembers(SimpleRemoteProxy, inspect.ismethod)
+        if not name.startswith("_") and callable(method)
+    }
+
+    execenv.print(f"RemoteProxy methods: {len(remote_proxy_methods)}")
+    execenv.print(f"SimpleRemoteProxy methods: {len(simple_proxy_methods)}")
+
+    # Methods in SimpleRemoteProxy but not in RemoteProxy
+    extra_in_simple = simple_proxy_methods - remote_proxy_methods
+    if extra_in_simple:
+        execenv.print(
+            "\n🔍 Methods in SimpleRemoteProxy but not in RemoteProxy "
+            "(expected, as SimpleRemoteProxy is a subset):"
+        )
+        for method in sorted(extra_in_simple):
+            execenv.print(f"  - {method}")
+
+    # Methods in RemoteProxy but not in SimpleRemoteProxy
+    missing_in_simple = remote_proxy_methods - simple_proxy_methods
+    if missing_in_simple:
+        execenv.print(
+            "\n📝 Methods in RemoteProxy but not in SimpleRemoteProxy "
+            "(expected, as SimpleRemoteProxy is simpler):"
+        )
+        for method in sorted(missing_in_simple):
+            execenv.print(f"  - {method}")
+
+    # Common methods
+    common_methods = remote_proxy_methods & simple_proxy_methods
+    execenv.print(f"\n✅ Common methods: {len(common_methods)}")
+
+    execenv.print("✨ Client method comparison completed successfully!")
+
+
+def test_with_real_server() -> None:
+    """Run sigima.tests.common.client_unit_test suite with real DataLab server.
+
+    This function runs the comprehensive client unit test suite after launching
+    a real DataLab instance in the background, ensuring that the tests pass
+    with the real server.
+    """
+    execenv.print("\n=== Testing with Real DataLab Server ===")
+
+    # Launch DataLab application in the background
+    execenv.print("Launching DataLab in background...")
+    run_datalab_in_background(wait_until_ready=True)
+
+    # Import and run the comprehensive test from sigima
+    execenv.print("Running comprehensive client tests with real server...")
+    tester = RemoteClientTester()
+
+    # Initialize connection to real DataLab server using the existing port
+    if not tester.init_cdl():
+        raise ConnectionRefusedError(
+            "Failed to connect to DataLab server. "
+            "Make sure DataLab is running and accessible."
+        )
+
+    try:
+        # Run all tests
+        tester.run_comprehensive_test()
+        execenv.print("✨ All tests passed with real DataLab server!")
+    finally:
+        # Clean up
+        tester.close_datalab()
+
+
+if __name__ == "__main__":
+    test_compare_server_methods()
+    test_compare_client_methods()
+    test_with_real_server()
