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