Skip to content

Commit

Permalink
Add unit test for the connection file merging
Browse files Browse the repository at this point in the history
  • Loading branch information
jasongrout committed Jul 13, 2023
1 parent 54ec030 commit edf708e
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 7 deletions.
11 changes: 4 additions & 7 deletions ipykernel/kernelapp.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@

import atexit
import errno
import json
import logging
import os
import signal
Expand All @@ -25,7 +24,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
Expand All @@ -45,10 +43,11 @@
from traitlets.utils.importstring import import_item
from zmq.eventloop.zmqstream import ZMQStream

from .control import ControlThread
from .heartbeat import Heartbeat
from .connect import get_connection_info, write_connection_file

# local imports
from .control import ControlThread
from .heartbeat import Heartbeat
from .iostream import IOPubThread
from .ipkernel import IPythonKernel
from .parentpoller import ParentPollerUnix, ParentPollerWindows
Expand Down Expand Up @@ -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)
Expand Down
57 changes: 57 additions & 0 deletions ipykernel/tests/test_kernelapp.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,13 @@
from unittest.mock import patch

import pytest
from jupyter_core.paths import secure_write
from traitlets.config.loader import Config

from ipykernel.kernelapp import IPKernelApp

from .conftest import MockKernel
from .utils import TemporaryWorkingDirectory

try:
import trio
Expand Down Expand Up @@ -47,6 +50,60 @@ 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

app.close()
os.remove(cf)


@pytest.mark.skipif(trio is None, reason="requires trio")
def test_trio_loop():
app = IPKernelApp(trio_loop=True)
Expand Down

0 comments on commit edf708e

Please sign in to comment.