88 lines
2.6 KiB
Python
88 lines
2.6 KiB
Python
import asyncio
|
|
import unittest
|
|
|
|
from ipython_shell.jupyter_app import JupyterApp
|
|
from ipython_shell.jupyter_shell import JupyterShell
|
|
from ipython_shell.transport import InputRequest, KernelMessage
|
|
|
|
|
|
class JupyterShellTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_shell_keeps_namespace_between_cells(self):
|
|
shell = JupyterShell("persistent")
|
|
try:
|
|
first = [
|
|
message
|
|
async for message in shell.run_cell_stream(
|
|
"answer = 41; answer",
|
|
call_id="call-1",
|
|
)
|
|
]
|
|
second = [
|
|
message
|
|
async for message in shell.run_cell_stream(
|
|
"answer + 1",
|
|
call_id="call-2",
|
|
)
|
|
]
|
|
finally:
|
|
await shell.shutdown()
|
|
|
|
self.assertEqual(
|
|
next(
|
|
message.content["data"]["text/plain"]
|
|
for message in first
|
|
if isinstance(message, KernelMessage)
|
|
and message.msg_type == "execute_result"
|
|
),
|
|
"41",
|
|
)
|
|
self.assertEqual(
|
|
next(
|
|
message.content["data"]["text/plain"]
|
|
for message in second
|
|
if isinstance(message, KernelMessage)
|
|
and message.msg_type == "execute_result"
|
|
),
|
|
"42",
|
|
)
|
|
|
|
|
|
class JupyterAppTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_app_routes_input_reply_to_the_call(self):
|
|
app = JupyterApp()
|
|
stream = app.run_code_stream(
|
|
"answer = input('name? '); answer",
|
|
shell_name="new",
|
|
)
|
|
|
|
try:
|
|
while True:
|
|
event = await asyncio.wait_for(anext(stream), timeout=5)
|
|
if isinstance(event, InputRequest):
|
|
input_request = event
|
|
break
|
|
|
|
self.assertIsInstance(input_request, InputRequest)
|
|
assert isinstance(input_request, InputRequest)
|
|
self.assertTrue(input_request.call_id)
|
|
self.assertEqual(input_request.prompt, "name? ")
|
|
|
|
await app.reply_to_input(input_request.call_id, "Ada")
|
|
messages = [message async for message in stream]
|
|
finally:
|
|
await app.shutdown()
|
|
|
|
self.assertEqual(
|
|
next(
|
|
message.content["data"]["text/plain"]
|
|
for message in messages
|
|
if isinstance(message, KernelMessage)
|
|
and message.msg_type == "execute_result"
|
|
),
|
|
"'Ada'",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|