Add unit test for the connection file merging

ipykernel/kernelapp.py

@@ -25,7 +25,6 @@
 )
 from IPython.core.profiledir import ProfileDir
 from IPython.core.shellapp import InteractiveShellApp, shell_aliases, shell_flags
-from jupyter_client import write_connection_file
 from jupyter_client.connect import ConnectionFileMixin
 from jupyter_client.session import Session, session_aliases, session_flags
 from jupyter_core.paths import jupyter_runtime_dir
@@ -45,10 +44,10 @@
 from traitlets.utils.importstring import import_item
 from zmq.eventloop.zmqstream import ZMQStream
 
+# local imports
 from .control import ControlThread
 from .heartbeat import Heartbeat
-
-# local imports
+from .connect import get_connection_info, write_connection_file
 from .iostream import IOPubThread
 from .ipkernel import IPythonKernel
 from .parentpoller import ParentPollerUnix, ParentPollerWindows
@@ -275,9 +274,7 @@ def write_connection_file(self):
             # If the file exists, merge our info into it. For example, if the
             # original file had port number 0, we update with the actual port
             # used.
-            with open(cf, "r") as f:
-                existing_connection_info = json.load(f)
-
+            existing_connection_info = get_connection_info(cf, unpack=True)
             connection_info = dict(existing_connection_info, **connection_info)
             if connection_info == existing_connection_info:
                 self.log.debug("Connection file %s with current information already exists", cf)

ipykernel/tests/test_kernelapp.py

@@ -5,7 +5,9 @@
 
 import pytest
 
+from ipykernel import connect
 from ipykernel.kernelapp import IPKernelApp
+from jupyter_core.paths import secure_write
 
 from .conftest import MockKernel
 
@@ -47,6 +49,61 @@ def trigger_stop():
     app.close()
 
 
+def test_merge_connection_file():
+    cfg = Config()
+    with TemporaryWorkingDirectory() as d:
+        cfg.ProfileDir.location = d
+        cf = os.path.join(d, "kernel.json")
+        initial_connection_info = {
+            "ip": "1.2.3.4",
+            "transport": "tcp",
+            "shell_port": 0,
+            "hb_port": 0,
+            "iopub_port": 0,
+            "stdin_port": 0,
+            "control_port": 5,
+            "key": "abc123",
+            "signature_scheme": "hmac-sha256",
+            "kernel_name": "My Kernel"
+        }
+        # We cannot use connect.write_connection_file since
+        # it replaces port number 0 with a random port
+        # and we want IPKernelApp to do that replacement.
+        with secure_write(cf) as f:
+            json.dump(initial_connection_info, f)
+        assert os.path.exists(cf)
+
+        app = IPKernelApp(config=cfg, connection_file=cf)
+        app.initialize()
+
+        # Initialize should have merged the actual connection info
+        # with the connection info in the file
+        assert cf == app.abs_connection_file
+        assert os.path.exists(cf)
+
+        with open(cf) as f:
+            new_connection_info = json.load(f)
+
+        # ports originally set as 0 have been replaced
+        for port in ("shell", "hb", "iopub", "stdin"):
+            key = f"{port}_port"
+            # We initially had the port as 0
+            assert initial_connection_info[key] == 0
+            # the port is not 0 now
+            assert new_connection_info[key] > 0
+            # the port matches the port the kernel actually used
+            assert new_connection_info[key] == getattr(app, key)
+            del new_connection_info[key]
+            del initial_connection_info[key]
+
+        # everything else in the connection file is the same
+        assert initial_connection_info == new_connection_info
+
+        os.remove(self.connection_file)
+
+        app.cleanup_connection_file()
+        app.close()
+
 @pytest.mark.skipif(trio is None, reason="requires trio")
 def test_trio_loop():
     app = IPKernelApp(trio_loop=True)