Show More
@@ -1,203 +1,198 b'' | |||||
1 | """test the IPython Kernel""" |
|
1 | """test the IPython Kernel""" | |
2 |
|
2 | |||
3 | #------------------------------------------------------------------------------- |
|
3 | #------------------------------------------------------------------------------- | |
4 | # Copyright (C) 2013 The IPython Development Team |
|
4 | # Copyright (C) 2013 The IPython Development Team | |
5 | # |
|
5 | # | |
6 | # Distributed under the terms of the BSD License. The full license is in |
|
6 | # Distributed under the terms of the BSD License. The full license is in | |
7 | # the file COPYING, distributed as part of this software. |
|
7 | # the file COPYING, distributed as part of this software. | |
8 | #------------------------------------------------------------------------------- |
|
8 | #------------------------------------------------------------------------------- | |
9 |
|
9 | |||
10 | #------------------------------------------------------------------------------- |
|
10 | #------------------------------------------------------------------------------- | |
11 | # Imports |
|
11 | # Imports | |
12 | #------------------------------------------------------------------------------- |
|
12 | #------------------------------------------------------------------------------- | |
13 |
|
13 | |||
14 | import os |
|
14 | import os | |
15 | import shutil |
|
15 | import shutil | |
16 | import tempfile |
|
16 | import tempfile | |
17 |
|
17 | |||
18 | from Queue import Empty |
|
18 | from Queue import Empty | |
19 | from contextlib import contextmanager |
|
19 | from contextlib import contextmanager | |
20 | from subprocess import PIPE |
|
20 | from subprocess import PIPE | |
21 |
|
21 | |||
22 | import nose.tools as nt |
|
22 | import nose.tools as nt | |
23 |
|
23 | |||
24 | from IPython.zmq.blockingkernelmanager import BlockingKernelManager |
|
24 | from IPython.zmq.blockingkernelmanager import BlockingKernelManager | |
25 | from IPython.zmq.tests.test_message_spec import execute, flush_channels |
|
25 | from IPython.zmq.tests.test_message_spec import execute, flush_channels | |
26 | from IPython.testing import decorators as dec |
|
26 | from IPython.testing import decorators as dec | |
27 | from IPython.utils import path, py3compat |
|
27 | from IPython.utils import path, py3compat | |
28 |
|
28 | |||
29 | #------------------------------------------------------------------------------- |
|
29 | #------------------------------------------------------------------------------- | |
30 | # Tests |
|
30 | # Tests | |
31 | #------------------------------------------------------------------------------- |
|
31 | #------------------------------------------------------------------------------- | |
32 |
|
32 | |||
33 | def setup(): |
|
33 | def setup(): | |
34 | """setup temporary IPYTHONDIR for tests""" |
|
34 | """setup temporary IPYTHONDIR for tests""" | |
35 | global IPYTHONDIR |
|
35 | global IPYTHONDIR | |
36 | global save_env |
|
36 | global save_env | |
37 | global save_get_ipython_dir |
|
37 | global save_get_ipython_dir | |
38 |
|
38 | |||
39 | IPYTHONDIR = tempfile.mkdtemp() |
|
39 | IPYTHONDIR = tempfile.mkdtemp() | |
40 |
|
40 | |||
41 | save_env = os.environ.copy() |
|
41 | save_env = os.environ.copy() | |
42 | os.environ["IPYTHONDIR"] = IPYTHONDIR |
|
42 | os.environ["IPYTHONDIR"] = IPYTHONDIR | |
43 |
|
43 | |||
44 | save_get_ipython_dir = path.get_ipython_dir |
|
44 | save_get_ipython_dir = path.get_ipython_dir | |
45 | path.get_ipython_dir = lambda : IPYTHONDIR |
|
45 | path.get_ipython_dir = lambda : IPYTHONDIR | |
46 |
|
46 | |||
47 |
|
47 | |||
48 | def teardown(): |
|
48 | def teardown(): | |
49 | path.get_ipython_dir = save_get_ipython_dir |
|
49 | path.get_ipython_dir = save_get_ipython_dir | |
50 | os.environ = save_env |
|
50 | os.environ = save_env | |
51 |
|
51 | |||
52 | try: |
|
52 | try: | |
53 | shutil.rmtree(IPYTHONDIR) |
|
53 | shutil.rmtree(IPYTHONDIR) | |
54 | except (OSError, IOError): |
|
54 | except (OSError, IOError): | |
55 | # no such file |
|
55 | # no such file | |
56 | pass |
|
56 | pass | |
57 |
|
57 | |||
58 |
|
58 | |||
59 | @contextmanager |
|
59 | @contextmanager | |
60 | def new_kernel(): |
|
60 | def new_kernel(): | |
61 | """start a kernel in a subprocess, and wait for it to be ready |
|
61 | """start a kernel in a subprocess, and wait for it to be ready | |
62 |
|
62 | |||
63 | Returns |
|
63 | Returns | |
64 | ------- |
|
64 | ------- | |
65 | kernel_manager: connected KernelManager instance |
|
65 | kernel_manager: connected KernelManager instance | |
66 | """ |
|
66 | """ | |
67 | KM = BlockingKernelManager() |
|
67 | KM = BlockingKernelManager() | |
68 |
|
68 | |||
69 | KM.start_kernel(stdout=PIPE, stderr=PIPE) |
|
69 | KM.start_kernel(stdout=PIPE, stderr=PIPE) | |
70 | KM.start_channels() |
|
70 | KM.start_channels() | |
71 |
|
71 | |||
72 | # wait for kernel to be ready |
|
72 | # wait for kernel to be ready | |
73 | KM.shell_channel.execute("import sys") |
|
73 | KM.shell_channel.execute("import sys") | |
74 | KM.shell_channel.get_msg(block=True, timeout=5) |
|
74 | KM.shell_channel.get_msg(block=True, timeout=5) | |
75 | flush_channels(KM) |
|
75 | flush_channels(KM) | |
76 | try: |
|
76 | try: | |
77 | yield KM |
|
77 | yield KM | |
78 | finally: |
|
78 | finally: | |
79 | KM.stop_channels() |
|
79 | KM.stop_channels() | |
80 | KM.shutdown_kernel() |
|
80 | KM.shutdown_kernel() | |
81 |
|
81 | |||
82 |
|
82 | |||
83 | def assemble_output(iopub): |
|
83 | def assemble_output(iopub): | |
84 | """assemble stdout/err from an execution""" |
|
84 | """assemble stdout/err from an execution""" | |
85 | stdout = '' |
|
85 | stdout = '' | |
86 | stderr = '' |
|
86 | stderr = '' | |
87 | while True: |
|
87 | while True: | |
88 | msg = iopub.get_msg(block=True, timeout=1) |
|
88 | msg = iopub.get_msg(block=True, timeout=1) | |
89 | msg_type = msg['msg_type'] |
|
89 | msg_type = msg['msg_type'] | |
90 | content = msg['content'] |
|
90 | content = msg['content'] | |
91 | if msg_type == 'status' and content['execution_state'] == 'idle': |
|
91 | if msg_type == 'status' and content['execution_state'] == 'idle': | |
92 | # idle message signals end of output |
|
92 | # idle message signals end of output | |
93 | break |
|
93 | break | |
94 | elif msg['msg_type'] == 'stream': |
|
94 | elif msg['msg_type'] == 'stream': | |
95 | if content['name'] == 'stdout': |
|
95 | if content['name'] == 'stdout': | |
96 | stdout = stdout + content['data'] |
|
96 | stdout = stdout + content['data'] | |
97 | elif content['name'] == 'stderr': |
|
97 | elif content['name'] == 'stderr': | |
98 | stderr = stderr + content['data'] |
|
98 | stderr = stderr + content['data'] | |
99 | else: |
|
99 | else: | |
100 | raise KeyError("bad stream: %r" % content['name']) |
|
100 | raise KeyError("bad stream: %r" % content['name']) | |
101 | else: |
|
101 | else: | |
102 | # other output, ignored |
|
102 | # other output, ignored | |
103 | pass |
|
103 | pass | |
104 | return stdout, stderr |
|
104 | return stdout, stderr | |
105 |
|
105 | |||
106 |
|
106 | |||
107 | def _check_mp_mode(km, expected=False, stream="stdout"): |
|
107 | def _check_mp_mode(km, expected=False, stream="stdout"): | |
108 | execute(km=km, code="import sys") |
|
108 | execute(km=km, code="import sys") | |
109 | flush_channels(km) |
|
109 | flush_channels(km) | |
110 | msg_id, content = execute(km=km, code="print (sys.%s._check_mp_mode())" % stream) |
|
110 | msg_id, content = execute(km=km, code="print (sys.%s._check_mp_mode())" % stream) | |
111 | stdout, stderr = assemble_output(km.iopub_channel) |
|
111 | stdout, stderr = assemble_output(km.iopub_channel) | |
112 | nt.assert_equal(eval(stdout.strip()), expected) |
|
112 | nt.assert_equal(eval(stdout.strip()), expected) | |
113 |
|
113 | |||
114 |
|
114 | |||
115 | def test_simple_print(): |
|
115 | def test_simple_print(): | |
116 | """simple print statement in kernel""" |
|
116 | """simple print statement in kernel""" | |
117 | with new_kernel() as km: |
|
117 | with new_kernel() as km: | |
118 | iopub = km.iopub_channel |
|
118 | iopub = km.iopub_channel | |
119 | msg_id, content = execute(km=km, code="print ('hi')") |
|
119 | msg_id, content = execute(km=km, code="print ('hi')") | |
120 | stdout, stderr = assemble_output(iopub) |
|
120 | stdout, stderr = assemble_output(iopub) | |
121 | nt.assert_equal(stdout, 'hi\n') |
|
121 | nt.assert_equal(stdout, 'hi\n') | |
122 | nt.assert_equal(stderr, '') |
|
122 | nt.assert_equal(stderr, '') | |
123 | _check_mp_mode(km, expected=False) |
|
123 | _check_mp_mode(km, expected=False) | |
124 | print ('hello') |
|
124 | print ('hello') | |
125 |
|
125 | |||
126 |
|
126 | |||
127 | def test_subprocess_print(): |
|
127 | def test_subprocess_print(): | |
128 | """printing from forked mp.Process""" |
|
128 | """printing from forked mp.Process""" | |
129 | with new_kernel() as km: |
|
129 | with new_kernel() as km: | |
130 | iopub = km.iopub_channel |
|
130 | iopub = km.iopub_channel | |
131 |
|
131 | |||
132 | _check_mp_mode(km, expected=False) |
|
132 | _check_mp_mode(km, expected=False) | |
133 | flush_channels(km) |
|
133 | flush_channels(km) | |
134 | np = 5 |
|
134 | np = 5 | |
135 | code = '\n'.join([ |
|
135 | code = '\n'.join([ | |
|
136 | "from __future__ import print_function", | |||
136 | "import multiprocessing as mp", |
|
137 | "import multiprocessing as mp", | |
137 | "def f(x):", |
|
138 | "pool = [mp.Process(target=print, args=('hello', i,)) for i in range(%i)]" % np, | |
138 | " print('hello',x)", |
|
|||
139 | "pool = [mp.Process(target=f,args=(i,)) for i in range(%i)]" % np, |
|
|||
140 | "for p in pool: p.start()", |
|
139 | "for p in pool: p.start()", | |
141 | "for p in pool: p.join()" |
|
140 | "for p in pool: p.join()" | |
142 | ]) |
|
141 | ]) | |
143 |
|
142 | |||
144 | expected = '\n'.join([ |
|
143 | expected = '\n'.join([ | |
145 | "hello %s" % i for i in range(np) |
|
144 | "hello %s" % i for i in range(np) | |
146 | ]) + '\n' |
|
145 | ]) + '\n' | |
147 |
|
146 | |||
148 | msg_id, content = execute(km=km, code=code) |
|
147 | msg_id, content = execute(km=km, code=code) | |
149 | stdout, stderr = assemble_output(iopub) |
|
148 | stdout, stderr = assemble_output(iopub) | |
150 | nt.assert_equal(stdout.count("hello"), np, stdout) |
|
149 | nt.assert_equal(stdout.count("hello"), np, stdout) | |
151 | for n in range(np): |
|
150 | for n in range(np): | |
152 | nt.assert_equal(stdout.count(str(n)), 1, stdout) |
|
151 | nt.assert_equal(stdout.count(str(n)), 1, stdout) | |
153 | nt.assert_equal(stderr, '') |
|
152 | nt.assert_equal(stderr, '') | |
154 | _check_mp_mode(km, expected=False) |
|
153 | _check_mp_mode(km, expected=False) | |
155 | _check_mp_mode(km, expected=False, stream="stderr") |
|
154 | _check_mp_mode(km, expected=False, stream="stderr") | |
156 |
|
155 | |||
157 |
|
156 | |||
158 | def test_subprocess_noprint(): |
|
157 | def test_subprocess_noprint(): | |
159 | """mp.Process without print doesn't trigger iostream mp_mode""" |
|
158 | """mp.Process without print doesn't trigger iostream mp_mode""" | |
160 | with new_kernel() as km: |
|
159 | with new_kernel() as km: | |
161 | iopub = km.iopub_channel |
|
160 | iopub = km.iopub_channel | |
162 |
|
161 | |||
163 | np = 5 |
|
162 | np = 5 | |
164 | code = '\n'.join([ |
|
163 | code = '\n'.join([ | |
165 | "import multiprocessing as mp", |
|
164 | "import multiprocessing as mp", | |
166 | "def f(x):", |
|
165 | "pool = [mp.Process(target=range,args=(i,)) for i in range(%i)]" % np, | |
167 | " return x", |
|
|||
168 | "pool = [mp.Process(target=f,args=(i,)) for i in range(%i)]" % np, |
|
|||
169 | "for p in pool: p.start()", |
|
166 | "for p in pool: p.start()", | |
170 | "for p in pool: p.join()" |
|
167 | "for p in pool: p.join()" | |
171 | ]) |
|
168 | ]) | |
172 |
|
169 | |||
173 | msg_id, content = execute(km=km, code=code) |
|
170 | msg_id, content = execute(km=km, code=code) | |
174 | stdout, stderr = assemble_output(iopub) |
|
171 | stdout, stderr = assemble_output(iopub) | |
175 | nt.assert_equal(stdout, '') |
|
172 | nt.assert_equal(stdout, '') | |
176 | nt.assert_equal(stderr, '') |
|
173 | nt.assert_equal(stderr, '') | |
177 |
|
174 | |||
178 | _check_mp_mode(km, expected=False) |
|
175 | _check_mp_mode(km, expected=False) | |
179 | _check_mp_mode(km, expected=False, stream="stderr") |
|
176 | _check_mp_mode(km, expected=False, stream="stderr") | |
180 |
|
177 | |||
181 |
|
178 | |||
182 | def test_subprocess_error(): |
|
179 | def test_subprocess_error(): | |
183 | """error in mp.Process doesn't crash""" |
|
180 | """error in mp.Process doesn't crash""" | |
184 | with new_kernel() as km: |
|
181 | with new_kernel() as km: | |
185 | iopub = km.iopub_channel |
|
182 | iopub = km.iopub_channel | |
186 |
|
183 | |||
187 | code = '\n'.join([ |
|
184 | code = '\n'.join([ | |
188 | "import multiprocessing as mp", |
|
185 | "import multiprocessing as mp", | |
189 | "def f():", |
|
186 | "p = mp.Process(target=int, args=('hi',))", | |
190 | " return 1/0", |
|
|||
191 | "p = mp.Process(target=f)", |
|
|||
192 | "p.start()", |
|
187 | "p.start()", | |
193 | "p.join()", |
|
188 | "p.join()", | |
194 | ]) |
|
189 | ]) | |
195 |
|
190 | |||
196 | msg_id, content = execute(km=km, code=code) |
|
191 | msg_id, content = execute(km=km, code=code) | |
197 | stdout, stderr = assemble_output(iopub) |
|
192 | stdout, stderr = assemble_output(iopub) | |
198 | nt.assert_equal(stdout, '') |
|
193 | nt.assert_equal(stdout, '') | |
199 |
nt.assert_true(" |
|
194 | nt.assert_true("ValueError" in stderr, stderr) | |
200 |
|
195 | |||
201 | _check_mp_mode(km, expected=False) |
|
196 | _check_mp_mode(km, expected=False) | |
202 | _check_mp_mode(km, expected=False, stream="stderr") |
|
197 | _check_mp_mode(km, expected=False, stream="stderr") | |
203 |
|
198 |
General Comments 0
You need to be logged in to leave comments.
Login now