mirror of
https://github.com/aykhans/AzSuicideDataVisualization.git
synced 2025-04-22 10:28:02 +00:00
257 lines
7.4 KiB
Python
257 lines
7.4 KiB
Python
"""Testing utils for jupyter_client tests
|
|
|
|
"""
|
|
import json
|
|
import os
|
|
import sys
|
|
from tempfile import TemporaryDirectory
|
|
from typing import Dict
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from jupyter_client import AsyncKernelManager
|
|
from jupyter_client import AsyncMultiKernelManager
|
|
from jupyter_client import KernelManager
|
|
from jupyter_client import MultiKernelManager
|
|
|
|
pjoin = os.path.join
|
|
|
|
skip_win32 = pytest.mark.skipif(sys.platform.startswith("win"), reason="Windows")
|
|
|
|
|
|
sample_kernel_json = {
|
|
"argv": ["cat", "{connection_file}"],
|
|
"display_name": "Test kernel",
|
|
}
|
|
|
|
|
|
def install_kernel(kernels_dir, argv=None, name="test", display_name=None):
|
|
"""install a kernel in a kernels directory"""
|
|
kernel_dir = pjoin(kernels_dir, name)
|
|
os.makedirs(kernel_dir)
|
|
kernel_json = {
|
|
"argv": argv or sample_kernel_json["argv"],
|
|
"display_name": display_name or sample_kernel_json["display_name"],
|
|
}
|
|
json_file = pjoin(kernel_dir, "kernel.json")
|
|
with open(json_file, "w") as f:
|
|
json.dump(kernel_json, f)
|
|
return kernel_dir
|
|
|
|
|
|
class test_env(object):
|
|
"""Set Jupyter path variables to a temporary directory
|
|
|
|
Useful as a context manager or with explicit start/stop
|
|
"""
|
|
|
|
def start(self):
|
|
self.test_dir = td = TemporaryDirectory()
|
|
self.env_patch = patch.dict(
|
|
os.environ,
|
|
{
|
|
"JUPYTER_CONFIG_DIR": pjoin(td.name, "jupyter"),
|
|
"JUPYTER_DATA_DIR": pjoin(td.name, "jupyter_data"),
|
|
"JUPYTER_RUNTIME_DIR": pjoin(td.name, "jupyter_runtime"),
|
|
"IPYTHONDIR": pjoin(td.name, "ipython"),
|
|
"TEST_VARS": "test_var_1",
|
|
},
|
|
)
|
|
self.env_patch.start()
|
|
|
|
def stop(self):
|
|
self.env_patch.stop()
|
|
try:
|
|
self.test_dir.cleanup()
|
|
except (PermissionError, NotADirectoryError):
|
|
if os.name != 'nt':
|
|
raise
|
|
|
|
def __enter__(self):
|
|
self.start()
|
|
return self.test_dir.name
|
|
|
|
def __exit__(self, *exc_info):
|
|
self.stop()
|
|
|
|
|
|
def execute(code="", kc=None, **kwargs):
|
|
"""wrapper for doing common steps for validating an execution request"""
|
|
from .test_message_spec import validate_message
|
|
|
|
if kc is None:
|
|
kc = KC # noqa
|
|
msg_id = kc.execute(code=code, **kwargs)
|
|
reply = kc.get_shell_msg(timeout=TIMEOUT) # noqa
|
|
validate_message(reply, "execute_reply", msg_id)
|
|
busy = kc.get_iopub_msg(timeout=TIMEOUT) # noqa
|
|
validate_message(busy, "status", msg_id)
|
|
assert busy["content"]["execution_state"] == "busy"
|
|
|
|
if not kwargs.get("silent"):
|
|
execute_input = kc.get_iopub_msg(timeout=TIMEOUT) # noqa
|
|
validate_message(execute_input, "execute_input", msg_id)
|
|
assert execute_input["content"]["code"] == code
|
|
|
|
return msg_id, reply["content"]
|
|
|
|
|
|
class RecordCallMixin:
|
|
method_calls: Dict[str, int]
|
|
|
|
def __init__(self, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.method_calls = {}
|
|
|
|
def record(self, method_name: str) -> None:
|
|
if method_name not in self.method_calls:
|
|
self.method_calls[method_name] = 0
|
|
self.method_calls[method_name] += 1
|
|
|
|
def call_count(self, method_name: str) -> int:
|
|
if method_name not in self.method_calls:
|
|
self.method_calls[method_name] = 0
|
|
return self.method_calls[method_name]
|
|
|
|
def reset_counts(self) -> None:
|
|
for record in self.method_calls:
|
|
self.method_calls[record] = 0
|
|
|
|
|
|
def subclass_recorder(f):
|
|
def wrapped(self, *args, **kwargs):
|
|
# record this call
|
|
self.record(f.__name__)
|
|
method = getattr(self._superclass, f.__name__)
|
|
# call the superclass method
|
|
r = method(self, *args, **kwargs)
|
|
# call anything defined in the actual class method
|
|
f(self, *args, **kwargs)
|
|
return r
|
|
|
|
return wrapped
|
|
|
|
|
|
class KMSubclass(RecordCallMixin):
|
|
@subclass_recorder
|
|
def start_kernel(self, **kw):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def shutdown_kernel(self, now=False, restart=False):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def restart_kernel(self, now=False, **kw):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def interrupt_kernel(self):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def request_shutdown(self, restart=False):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def finish_shutdown(self, waittime=None, pollinterval=0.1, restart=False):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def _launch_kernel(self, kernel_cmd, **kw):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def _kill_kernel(self):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def cleanup_resources(self, restart=False):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def signal_kernel(self, signum: int):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def is_alive(self):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def _send_kernel_sigterm(self, restart: bool = False):
|
|
"""Record call and defer to superclass"""
|
|
|
|
|
|
class SyncKMSubclass(KMSubclass, KernelManager):
|
|
"""Used to test subclass hierarchies to ensure methods are called when expected."""
|
|
|
|
_superclass = KernelManager
|
|
|
|
|
|
class AsyncKMSubclass(KMSubclass, AsyncKernelManager):
|
|
"""Used to test subclass hierarchies to ensure methods are called when expected."""
|
|
|
|
_superclass = AsyncKernelManager
|
|
|
|
|
|
class MKMSubclass(RecordCallMixin):
|
|
def _kernel_manager_class_default(self):
|
|
return "jupyter_client.tests.utils.SyncKMSubclass"
|
|
|
|
@subclass_recorder
|
|
def get_kernel(self, kernel_id):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def remove_kernel(self, kernel_id):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def start_kernel(self, kernel_name=None, **kwargs):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def shutdown_kernel(self, kernel_id, now=False, restart=False):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def restart_kernel(self, kernel_id, now=False):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def interrupt_kernel(self, kernel_id):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def request_shutdown(self, kernel_id, restart=False):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def finish_shutdown(self, kernel_id, waittime=None, pollinterval=0.1, restart=False):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def cleanup_resources(self, kernel_id, restart=False):
|
|
"""Record call and defer to superclass"""
|
|
|
|
@subclass_recorder
|
|
def shutdown_all(self, now=False):
|
|
"""Record call and defer to superclass"""
|
|
|
|
|
|
class SyncMKMSubclass(MKMSubclass, MultiKernelManager):
|
|
|
|
_superclass = MultiKernelManager
|
|
|
|
def _kernel_manager_class_default(self):
|
|
return "jupyter_client.tests.utils.SyncKMSubclass"
|
|
|
|
|
|
class AsyncMKMSubclass(MKMSubclass, AsyncMultiKernelManager):
|
|
|
|
_superclass = AsyncMultiKernelManager
|
|
|
|
def _kernel_manager_class_default(self):
|
|
return "jupyter_client.tests.utils.AsyncKMSubclass"
|