test_ipython.py

from __future__ import print_function, division, absolute_import

import mock

import pytest
from toolz import first
import tornado

from distributed import Client
from distributed.utils_test import cluster, mock_ipython
from distributed.utils_test import loop, zmq_ctx  # noqa F401


def need_functional_ipython(func):
    try:
        import ipykernel  # noqa: F401
        import jupyter_client  # noqa: F401
    except ImportError:
        return pytest.mark.skip("need ipykernel and jupyter_client installed")(func)
    if tornado.version_info >= (5,):
        # https://github.com/ipython/ipykernel/issues/277
        return pytest.mark.skip("IPython kernel broken with Tornado 5")(func)
    else:
        return func


@pytest.mark.ipython
@need_functional_ipython
def test_start_ipython_workers(loop, zmq_ctx):
    from jupyter_client import BlockingKernelClient

    with cluster(1) as (s, [a]):
        with Client(s["address"], loop=loop) as e:
            info_dict = e.start_ipython_workers()
            info = first(info_dict.values())
            key = info.pop("key")
            kc = BlockingKernelClient(**info)
            kc.session.key = key
            kc.start_channels()
            kc.wait_for_ready(timeout=10)
            msg_id = kc.execute("worker")
            reply = kc.get_shell_msg(timeout=10)
            assert reply["parent_header"]["msg_id"] == msg_id
            assert reply["content"]["status"] == "ok"
            kc.stop_channels()


@pytest.mark.ipython
@need_functional_ipython
def test_start_ipython_scheduler(loop, zmq_ctx):
    from jupyter_client import BlockingKernelClient

    with cluster(1) as (s, [a]):
        with Client(s["address"], loop=loop) as e:
            info = e.start_ipython_scheduler()
            key = info.pop("key")
            kc = BlockingKernelClient(**info)
            kc.session.key = key
            kc.start_channels()
            msg_id = kc.execute("scheduler")
            reply = kc.get_shell_msg(timeout=10)
            kc.stop_channels()


@pytest.mark.ipython
@need_functional_ipython
def test_start_ipython_scheduler_magic(loop, zmq_ctx):
    with cluster(1) as (s, [a]):
        with Client(s["address"], loop=loop) as e, mock_ipython() as ip:
            info = e.start_ipython_scheduler()

        expected = [
            {"magic_kind": "line", "magic_name": "scheduler"},
            {"magic_kind": "cell", "magic_name": "scheduler"},
        ]

        call_kwargs_list = [
            kwargs for (args, kwargs) in ip.register_magic_function.call_args_list
        ]
        assert call_kwargs_list == expected
        magic = ip.register_magic_function.call_args_list[0][0][0]
        magic(line="", cell="scheduler")


@pytest.mark.ipython
@need_functional_ipython
def test_start_ipython_workers_magic(loop, zmq_ctx):
    with cluster(2) as (s, [a, b]):

        with Client(s["address"], loop=loop) as e, mock_ipython() as ip:
            workers = list(e.nthreads())[:2]
            names = ["magic%i" % i for i in range(len(workers))]
            info_dict = e.start_ipython_workers(workers, magic_names=names)

        expected = [
            {"magic_kind": "line", "magic_name": "remote"},
            {"magic_kind": "cell", "magic_name": "remote"},
            {"magic_kind": "line", "magic_name": "magic0"},
            {"magic_kind": "cell", "magic_name": "magic0"},
            {"magic_kind": "line", "magic_name": "magic1"},
            {"magic_kind": "cell", "magic_name": "magic1"},
        ]
        call_kwargs_list = [
            kwargs for (args, kwargs) in ip.register_magic_function.call_args_list
        ]
        assert call_kwargs_list == expected
        assert ip.register_magic_function.call_count == 6
        magics = [args[0][0] for args in ip.register_magic_function.call_args_list[2:]]
        magics[-1](line="", cell="worker")
        [m.client.stop_channels() for m in magics]


@pytest.mark.ipython
@need_functional_ipython
def test_start_ipython_workers_magic_asterix(loop, zmq_ctx):
    with cluster(2) as (s, [a, b]):

        with Client(s["address"], loop=loop) as e, mock_ipython() as ip:
            workers = list(e.nthreads())[:2]
            info_dict = e.start_ipython_workers(workers, magic_names="magic_*")

        expected = [
            {"magic_kind": "line", "magic_name": "remote"},
            {"magic_kind": "cell", "magic_name": "remote"},
            {"magic_kind": "line", "magic_name": "magic_0"},
            {"magic_kind": "cell", "magic_name": "magic_0"},
            {"magic_kind": "line", "magic_name": "magic_1"},
            {"magic_kind": "cell", "magic_name": "magic_1"},
        ]
        call_kwargs_list = [
            kwargs for (args, kwargs) in ip.register_magic_function.call_args_list
        ]
        assert call_kwargs_list == expected
        assert ip.register_magic_function.call_count == 6
        magics = [args[0][0] for args in ip.register_magic_function.call_args_list[2:]]
        magics[-1](line="", cell="worker")
        [m.client.stop_channels() for m in magics]


@pytest.mark.ipython
@need_functional_ipython
def test_start_ipython_remote(loop, zmq_ctx):
    from distributed._ipython_utils import remote_magic

    with cluster(1) as (s, [a]):
        with Client(s["address"], loop=loop) as e, mock_ipython() as ip:
            worker = first(e.nthreads())
            ip.user_ns["info"] = e.start_ipython_workers(worker)[worker]
            remote_magic("info 1")  # line magic
            remote_magic("info", "worker")  # cell magic

        expected = [
            ((remote_magic,), {"magic_kind": "line", "magic_name": "remote"}),
            ((remote_magic,), {"magic_kind": "cell", "magic_name": "remote"}),
        ]
        assert ip.register_magic_function.call_args_list == expected
        assert ip.register_magic_function.call_count == 2


@pytest.mark.ipython
@need_functional_ipython
def test_start_ipython_qtconsole(loop):
    Popen = mock.Mock()
    with cluster() as (s, [a, b]):
        with mock.patch("distributed._ipython_utils.Popen", Popen), Client(
            s["address"], loop=loop
        ) as e:
            worker = first(e.nthreads())
            e.start_ipython_workers(worker, qtconsole=True)
            e.start_ipython_workers(worker, qtconsole=True, qtconsole_args=["--debug"])
    assert Popen.call_count == 2
    (cmd,), kwargs = Popen.call_args_list[0]
    assert cmd[:3] == ["jupyter", "qtconsole", "--existing"]
    (cmd,), kwargs = Popen.call_args_list[1]
    assert cmd[-1:] == ["--debug"]