1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
|
import os
import time
from queue import Empty
import unittest
from flaky import flaky
from qtconsole.manager import QtKernelManager
class Tests(unittest.TestCase):
def setUp(self):
"""Open a kernel."""
# Prevent tests assertions related with message type from failing
# due to a debug warning
os.environ['PYDEVD_DISABLE_FILE_VALIDATION'] = '1'
self.kernel_manager = QtKernelManager()
self.kernel_manager.start_kernel()
self.kernel_client = self.kernel_manager.client()
self.kernel_client.start_channels(shell=True, iopub=True)
self.blocking_client = self.kernel_client.blocking_client()
self.blocking_client.start_channels(shell=True, iopub=True)
self.comm_manager = self.kernel_client.comm_manager
# Check if client is working
self.blocking_client.execute('print(0)')
try:
self._get_next_msg()
self._get_next_msg()
except TimeoutError:
# Maybe it works now?
self.blocking_client.execute('print(0)')
self._get_next_msg()
self._get_next_msg()
def tearDown(self):
"""Close the kernel."""
if self.kernel_manager:
self.kernel_manager.shutdown_kernel(now=True)
if self.kernel_client:
self.kernel_client.shutdown()
def _get_next_msg(self, timeout=10):
# Get status messages
timeout_time = time.time() + timeout
msg_type = 'status'
while msg_type == 'status':
if timeout_time < time.time():
raise TimeoutError
try:
msg = self.blocking_client.get_iopub_msg(timeout=3)
msg_type = msg['header']['msg_type']
except Empty:
pass
return msg
@flaky(max_runs=10)
def test_kernel_to_frontend(self):
"""Communicate from the kernel to the frontend."""
comm_manager = self.comm_manager
blocking_client = self.blocking_client
class DummyCommHandler():
def __init__(self):
comm_manager.register_target('test_api', self.comm_open)
self.last_msg = None
def comm_open(self, comm, msg):
comm.on_msg(self.comm_message)
comm.on_close(self.comm_message)
self.last_msg = msg['content']['data']
self.comm = comm
def comm_message(self, msg):
self.last_msg = msg['content']['data']
handler = DummyCommHandler()
blocking_client.execute(
"from ipykernel.comm import Comm\n"
"comm = Comm(target_name='test_api', data='open')\n"
"comm.send('message')\n"
"comm.close('close')\n"
"del comm\n"
"print('Done')\n"
)
# Get input
msg = self._get_next_msg()
assert msg['header']['msg_type'] == 'execute_input'
# Open comm
msg = self._get_next_msg()
assert msg['header']['msg_type'] == 'comm_open'
comm_manager._dispatch(msg)
assert handler.last_msg == 'open'
assert handler.comm.comm_id == msg['content']['comm_id']
# Get message
msg = self._get_next_msg()
assert msg['header']['msg_type'] == 'comm_msg'
comm_manager._dispatch(msg)
assert handler.last_msg == 'message'
assert handler.comm.comm_id == msg['content']['comm_id']
# Get close
msg = self._get_next_msg()
assert msg['header']['msg_type'] == 'comm_close'
comm_manager._dispatch(msg)
assert handler.last_msg == 'close'
assert handler.comm.comm_id == msg['content']['comm_id']
# Get close
msg = self._get_next_msg()
assert msg['header']['msg_type'] == 'stream'
@flaky(max_runs=10)
def test_frontend_to_kernel(self):
"""Communicate from the frontend to the kernel."""
comm_manager = self.comm_manager
blocking_client = self.blocking_client
blocking_client.execute(
"class DummyCommHandler():\n"
" def __init__(self):\n"
" get_ipython().kernel.comm_manager.register_target(\n"
" 'test_api', self.comm_open)\n"
" def comm_open(self, comm, msg):\n"
" comm.on_msg(self.comm_message)\n"
" comm.on_close(self.comm_message)\n"
" print(msg['content']['data'])\n"
" def comm_message(self, msg):\n"
" print(msg['content']['data'])\n"
"dummy = DummyCommHandler()\n"
)
# Get input
msg = self._get_next_msg()
assert msg['header']['msg_type'] == 'execute_input'
# Open comm
comm = comm_manager.new_comm('test_api', data='open')
msg = self._get_next_msg()
assert msg['header']['msg_type'] == 'stream'
assert msg['content']['text'] == 'open\n'
# Get message
comm.send('message')
msg = self._get_next_msg()
assert msg['header']['msg_type'] == 'stream'
assert msg['content']['text'] == 'message\n'
# Get close
comm.close('close')
msg = self._get_next_msg()
# Received message has a header and parent header. The parent header has
# the info about the close message type in Python 3
assert msg['parent_header']['msg_type'] == 'comm_close'
assert msg['msg_type'] == 'stream'
assert msg['content']['text'] == 'close\n'
if __name__ == "__main__":
unittest.main()
|
