|
| 1 | +# Copyright (c) Jupyter Development Team. |
| 2 | +# Distributed under the terms of the Modified BSD License. |
| 3 | + |
| 4 | +import json |
| 5 | +import os |
| 6 | +import sys |
| 7 | +from pathlib import Path |
| 8 | + |
| 9 | +import pytest |
| 10 | +from jupyter_core import paths |
| 11 | + |
| 12 | +try: |
| 13 | + import ipykernel # noqa |
| 14 | + from jupyter_client.manager import start_new_async_kernel |
| 15 | +except ImportError: |
| 16 | + import warnings |
| 17 | + |
| 18 | + warnings.warn( |
| 19 | + "The client plugin has not been installed. " |
| 20 | + "If you're trying to use this plugin and you've installed " |
| 21 | + "`pytest-jupyter`, there is likely one more step " |
| 22 | + "you need. Try: `pip install 'pytest-jupyter[client]'`" |
| 23 | + ) |
| 24 | + |
| 25 | +try: |
| 26 | + import resource |
| 27 | +except ImportError: |
| 28 | + # Windows |
| 29 | + resource = None # type: ignore |
| 30 | + |
| 31 | + |
| 32 | +# Handle resource limit |
| 33 | +# Ensure a minimal soft limit of DEFAULT_SOFT if the current hard limit is at least that much. |
| 34 | +if resource is not None: |
| 35 | + soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) |
| 36 | + |
| 37 | + DEFAULT_SOFT = 4096 |
| 38 | + if hard >= DEFAULT_SOFT: |
| 39 | + soft = DEFAULT_SOFT |
| 40 | + |
| 41 | + if hard < soft: |
| 42 | + hard = soft |
| 43 | + |
| 44 | + resource.setrlimit(resource.RLIMIT_NOFILE, (soft, hard)) |
| 45 | + |
| 46 | + |
| 47 | +@pytest.fixture |
| 48 | +def zmq_context(): |
| 49 | + import zmq |
| 50 | + |
| 51 | + ctx = zmq.asyncio.Context() |
| 52 | + yield ctx |
| 53 | + ctx.term() |
| 54 | + |
| 55 | + |
| 56 | +@pytest.fixture |
| 57 | +async def start_kernel(kernel_spec): |
| 58 | + km = None |
| 59 | + kc = None |
| 60 | + |
| 61 | + async def inner(kernel_name="echo", **kwargs): |
| 62 | + nonlocal km, kc |
| 63 | + km, kc = await start_new_async_kernel(kernel_name=kernel_name, **kwargs) |
| 64 | + return km, kc |
| 65 | + |
| 66 | + yield inner |
| 67 | + |
| 68 | + if kc and km: |
| 69 | + kc.stop_channels() |
| 70 | + await km.shutdown_kernel(now=True) |
| 71 | + assert km.context.closed |
| 72 | + km.context.destroy() |
| 73 | + km.context.term() |
| 74 | + |
| 75 | + |
| 76 | +@pytest.fixture() |
| 77 | +def kernel_dir(): |
| 78 | + return os.path.join(paths.jupyter_data_dir(), "kernels") |
| 79 | + |
| 80 | + |
| 81 | +@pytest.fixture |
| 82 | +def kernel_spec(kernel_dir): |
| 83 | + test_dir = Path(kernel_dir) / "echo" |
| 84 | + test_dir.mkdir(parents=True, exist_ok=True) |
| 85 | + argv = [sys.executable, "-m", "pytest_jupyter.echo_kernel", "-f", "{connection_file}"] |
| 86 | + kernel_data = {"argv": argv, "display_name": "echo", "language": "echo"} |
| 87 | + spec_file_path = Path(test_dir / "kernel.json") |
| 88 | + spec_file_path.write_text(json.dumps(kernel_data), "utf8") |
| 89 | + return str(test_dir) |
0 commit comments