async def test_jupyter_kernel_port_close(hass, caplog):
"""Test Jupyter kernel closing ports."""
sock, port_nums = await setup_script(hass, [dt(2020, 7, 1, 11, 0, 0, 0)], "")
#
# test the heartbeat loopback with some long and short messages
# also send messages to stdin and iopub, which ignore them
#
for i in range(5):
if i & 1:
msg = (f"hello {i} " * 40).encode("utf-8")
else:
msg = f"hello {i}".encode("utf-8")
await sock["hb_port"].send(msg)
await sock["iopub_port"].send(msg)
await sock["stdin_port"].send(msg)
assert await sock["hb_port"].recv() == msg
#
# now close and re-open each por
#
for name in PORT_NAMES:
sock[name].close()
kernel_reader, kernel_writer = await asyncio.open_connection("127.0.0.1", port_nums[name])
sock[name] = ZmqSocket(kernel_reader, kernel_writer, "ROUTER")
await sock[name].handshake()
#
# test the heartbeat loopback again
# also send messages to stdin and iopub, which ignore them
#
for i in range(5):
if i & 1:
msg = (f"hello {i} " * 40).encode("utf-8")
else:
msg = f"hello {i}".encode("utf-8")
await sock["hb_port"].send(msg)
await sock["iopub_port"].send(msg)
await sock["stdin_port"].send(msg)
assert await sock["hb_port"].recv() == msg
hass.bus.async_fire(EVENT_HOMEASSISTANT_STARTED)
#
# shut down the session via signature mismatch with bad key
#
await send(
sock["control_port"], "shutdown_request", {}, parent_header={}, identities={}, secret_key=b"bad_key"
)
#
# wait until the session ends, so the log receives the error message we check below
#
try:
await sock["iopub_port"].recv()
except EOFError:
pass
assert "signature mismatch: check_sig=" in caplog.text
async def setup_script(hass, now, source, no_connect=False):
"""Initialize and load the given pyscript."""
scripts = [
"/some/config/dir/pyscripts/hello.py",
]
with patch(
"custom_components.pyscript.os.path.isdir",
return_value=True), patch(
"custom_components.pyscript.glob.iglob",
return_value=scripts), patch(
"custom_components.pyscript.global_ctx.open",
mock_open(read_data=source),
create=True,
), patch(
"custom_components.pyscript.trigger.dt_now",
return_value=now), patch(
"homeassistant.config.load_yaml_config_file",
return_value={}), patch(
"custom_components.pyscript.install_requirements",
return_value=None,
):
assert await async_setup_component(hass, "pyscript", {DOMAIN: {}})
#
# I'm not sure how to run the mock all the time, so just force the dt_now()
# trigger function to return the given list of times in now.
#
def return_next_time():
nonlocal now
if isinstance(now, list):
if len(now) > 1:
return now.pop(0)
return now[0]
return now
trigger.__dict__["dt_now"] = return_next_time
kernel_state_var = "pyscript.jupyter_ports_1234"
kernel_cfg = {
"ip": "127.0.0.1",
"key": SECRET_KEY.decode("utf-8"),
"signature_scheme": "hmac-sha256",
"state_var": kernel_state_var,
}
if no_connect:
kernel_cfg["no_connect_timeout"] = 0.0
await hass.services.async_call("pyscript", "jupyter_kernel_start",
kernel_cfg)
while True:
ports_state = hass.states.get(kernel_state_var)
if ports_state is not None:
break
await asyncio.sleep(2e-3)
port_nums = json.loads(ports_state.state)
sock = {}
if no_connect:
return sock, port_nums
for name in PORT_NAMES:
kernel_reader, kernel_writer = await asyncio.open_connection(
"127.0.0.1", port_nums[name])
sock[name] = ZmqSocket(kernel_reader, kernel_writer, "ROUTER")
await sock[name].handshake()
return sock, port_nums
async def setup_script(hass, now, source, no_connect=False):
"""Initialize and load the given pyscript."""
conf_dir = hass.config.path(FOLDER)
file_contents = {f"{conf_dir}/hello.py": source}
mock_open = MockOpen()
for key, value in file_contents.items():
mock_open[key].read_data = value
def isfile_side_effect(arg):
return arg in file_contents
def glob_side_effect(path, recursive=None):
result = []
path_re = path.replace("*", "[^/]*").replace(".", "\\.")
path_re = path_re.replace("[^/]*[^/]*/", ".*")
for this_path in file_contents:
if re.match(path_re, this_path):
result.append(this_path)
return result
with patch(
"custom_components.pyscript.os.path.isdir", return_value=True
), patch("custom_components.pyscript.glob.iglob") as mock_glob, patch(
"custom_components.pyscript.global_ctx.open", mock_open
), patch(
"custom_components.pyscript.trigger.dt_now", return_value=now
), patch("custom_components.pyscript.open", mock_open), patch(
"homeassistant.config.load_yaml_config_file", return_value={}
), patch(
"custom_components.pyscript.install_requirements",
return_value=None,
), patch("custom_components.pyscript.watchdog_start",
return_value=None), patch(
"custom_components.pyscript.os.path.getmtime",
return_value=1000), patch(
"custom_components.pyscript.global_ctx.os.path.getmtime",
return_value=1000), patch(
"custom_components.pyscript.os.path.isfile"
) as mock_isfile:
mock_isfile.side_effect = isfile_side_effect
mock_glob.side_effect = glob_side_effect
assert await async_setup_component(hass, "pyscript", {DOMAIN: {}})
#
# I'm not sure how to run the mock all the time, so just force the dt_now()
# trigger function to return the given list of times in now.
#
def return_next_time():
nonlocal now
if isinstance(now, list):
if len(now) > 1:
return now.pop(0)
return now[0]
return now
trigger.__dict__["dt_now"] = return_next_time
kernel_state_var = "pyscript.jupyter_ports_1234"
kernel_cfg = {
"ip": "127.0.0.1",
"key": SECRET_KEY.decode("utf-8"),
"signature_scheme": "hmac-sha256",
"state_var": kernel_state_var,
}
if no_connect:
kernel_cfg["no_connect_timeout"] = 0.0
await hass.services.async_call("pyscript", "jupyter_kernel_start",
kernel_cfg)
while True:
ports_state = hass.states.get(kernel_state_var)
if ports_state is not None:
break
await asyncio.sleep(2e-3)
port_nums = json.loads(ports_state.state)
sock = {}
if no_connect:
return sock, port_nums
for name in PORT_NAMES:
kernel_reader, kernel_writer = await asyncio.open_connection(
"127.0.0.1", port_nums[name])
sock[name] = ZmqSocket(kernel_reader, kernel_writer, "ROUTER")
await sock[name].handshake()
return sock, port_nums