def _setup():
global reconnect_handler, session_manager, session_list, path, kernel_id, kernel_manager,\
individual_kernel_manager, response_id, good_msg, client, spark_events
# Mock kernel manager
client = MagicMock()
client.execute = MagicMock(return_value=response_id)
client.get_shell_msg = MagicMock(return_value=good_msg)
individual_kernel_manager = MagicMock()
individual_kernel_manager.client = MagicMock(return_value=client)
kernel_manager = MagicMock()
kernel_manager.get_kernel = MagicMock(
return_value=individual_kernel_manager)
# Mock session manager
session_list = [create_session_dict(path, kernel_id)]
session_manager = MagicMock()
session_manager.list_sessions = MagicMock(return_value=session_list)
#Mock spark events
spark_events = MagicMock()
# Create mocked reconnect_handler
ReconnectHandler.__bases__ = (SimpleObject, )
reconnect_handler = ReconnectHandler()
reconnect_handler.spark_events = spark_events
reconnect_handler.session_manager = session_manager
reconnect_handler.kernel_manager = kernel_manager
reconnect_handler.set_status = MagicMock()
reconnect_handler.finish = MagicMock()
reconnect_handler.get_body_argument = get_argument
reconnect_handler.current_user = '******'
def _setup():
global reconnect_handler, session_manager, session_list, path, kernel_id, kernel_manager,\
individual_kernel_manager, response_id, good_msg, client, spark_events, request
# Mock kernel manager
client = MagicMock()
client.execute = MagicMock(return_value=response_id)
client.get_shell_msg = MagicMock(return_value=good_msg)
individual_kernel_manager = MagicMock()
individual_kernel_manager.client = MagicMock(return_value=client)
kernel_manager = MagicMock()
kernel_manager.get_kernel = MagicMock(return_value=individual_kernel_manager)
# Mock session manager
session_list = [create_session_dict(path, kernel_id)]
session_manager = MagicMock()
session_manager.list_sessions = MagicMock(return_value=session_list)
# Mock spark events
spark_events = MagicMock()
# Mock request
request = MagicMock()
request.body = json.dumps({"path": path, "username": username, "password": password, "endpoint": endpoint})
# Create mocked reconnect_handler
ReconnectHandler.__bases__ = (SimpleObject,)
reconnect_handler = ReconnectHandler()
reconnect_handler.spark_events = spark_events
reconnect_handler.session_manager = session_manager
reconnect_handler.kernel_manager = kernel_manager
reconnect_handler.set_status = MagicMock()
reconnect_handler.finish = MagicMock()
reconnect_handler.current_user = '******'
reconnect_handler.request = request
def setUp(self):
# Mock kernel manager
self.client = MagicMock()
self.client.execute = MagicMock(return_value=self.response_id)
self.client.get_shell_msg = MagicMock(return_value=self.good_msg)
self.individual_kernel_manager = MagicMock()
self.individual_kernel_manager.client = MagicMock(
return_value=self.client)
self.kernel_manager = MagicMock()
self.kernel_manager.get_kernel = MagicMock(
return_value=self.individual_kernel_manager)
# Mock session manager
self.session_list = [
self.create_session_dict(self.path, self.kernel_id)
]
self.session_manager = MagicMock()
self.session_manager.list_sessions = MagicMock(
return_value=self.session_list)
self.session_manager.create_session = MagicMock(
return_value=self.create_session_dict(self.path, self.kernel_id))
# Mock spark events
self.spark_events = MagicMock()
# Mock request
self.request = MagicMock()
self.request.body = json.dumps({
"path": self.path,
"username": self.username,
"password": self.password,
"endpoint": self.endpoint,
"auth": self.auth
})
# Create mocked reconnect_handler
ReconnectHandler.__bases__ = (SimpleObject, )
self.reconnect_handler = ReconnectHandler()
self.reconnect_handler.spark_events = self.spark_events
self.reconnect_handler.session_manager = self.session_manager
self.reconnect_handler.kernel_manager = self.kernel_manager
self.reconnect_handler.set_status = MagicMock()
self.reconnect_handler.finish = MagicMock()
self.reconnect_handler.current_user = '******'
self.reconnect_handler.request = self.request
self.reconnect_handler.logger = MagicMock()
super(TestSparkMagicHandler, self).setUp()
def _setup():
global reconnect_handler, session_manager, session_list, path, kernel_id, kernel_manager,\
individual_kernel_manager, response_id, good_msg, client, spark_events, request
# Mock kernel manager
client = MagicMock()
client.execute = MagicMock(return_value=response_id)
client.get_shell_msg = MagicMock(return_value=good_msg)
individual_kernel_manager = MagicMock()
individual_kernel_manager.client = MagicMock(return_value=client)
kernel_manager = MagicMock()
kernel_manager.get_kernel = MagicMock(
return_value=individual_kernel_manager)
# Mock session manager
session_list = [create_session_dict(path, kernel_id)]
session_manager = MagicMock()
session_manager.list_sessions = MagicMock(return_value=session_list)
# Mock spark events
spark_events = MagicMock()
# Mock request
request = MagicMock()
request.body = json.dumps({
"path": path,
"username": username,
"password": password,
"endpoint": endpoint
})
# Create mocked reconnect_handler
ReconnectHandler.__bases__ = (SimpleObject, )
reconnect_handler = ReconnectHandler()
reconnect_handler.spark_events = spark_events
reconnect_handler.session_manager = session_manager
reconnect_handler.kernel_manager = kernel_manager
reconnect_handler.set_status = MagicMock()
reconnect_handler.finish = MagicMock()
reconnect_handler.current_user = '******'
reconnect_handler.request = request
class TestSparkMagicHandler(AsyncTestCase):
reconnect_handler = None
session_manager = None
kernel_manager = None
individual_kernel_manager = None
client = None
session_list = None
spark_events = None
path = 'some_path.ipynb'
kernel_id = '1'
kernel_name = 'pysparkkernel'
session_id = '1'
username = '******'
password = '******'
endpoint = 'http://endpoint.com'
auth = constants.AUTH_BASIC
response_id = '0'
good_msg = dict(content=dict(status='ok'))
bad_msg = dict(
content=dict(status='error', ename='SyntaxError', evalue='oh no!'))
request = None
def create_session_dict(self, path, kernel_id):
return dict(notebook=dict(path=path),
kernel=dict(id=kernel_id, name=self.kernel_name),
id=self.session_id)
def get_argument(self, key):
return dict(username=self.username,
password=self.password,
endpoint=self.endpoint,
path=self.path)[key]
def setUp(self):
# Mock kernel manager
self.client = MagicMock()
self.client.execute = MagicMock(return_value=self.response_id)
self.client.get_shell_msg = MagicMock(return_value=self.good_msg)
self.individual_kernel_manager = MagicMock()
self.individual_kernel_manager.client = MagicMock(
return_value=self.client)
self.kernel_manager = MagicMock()
self.kernel_manager.get_kernel = MagicMock(
return_value=self.individual_kernel_manager)
# Mock session manager
self.session_list = [
self.create_session_dict(self.path, self.kernel_id)
]
self.session_manager = MagicMock()
self.session_manager.list_sessions = MagicMock(
return_value=self.session_list)
self.session_manager.create_session = MagicMock(
return_value=self.create_session_dict(self.path, self.kernel_id))
# Mock spark events
self.spark_events = MagicMock()
# Mock request
self.request = MagicMock()
self.request.body = json.dumps({
"path": self.path,
"username": self.username,
"password": self.password,
"endpoint": self.endpoint,
"auth": self.auth
})
# Create mocked reconnect_handler
ReconnectHandler.__bases__ = (SimpleObject, )
self.reconnect_handler = ReconnectHandler()
self.reconnect_handler.spark_events = self.spark_events
self.reconnect_handler.session_manager = self.session_manager
self.reconnect_handler.kernel_manager = self.kernel_manager
self.reconnect_handler.set_status = MagicMock()
self.reconnect_handler.finish = MagicMock()
self.reconnect_handler.current_user = '******'
self.reconnect_handler.request = self.request
self.reconnect_handler.logger = MagicMock()
super(TestSparkMagicHandler, self).setUp()
def test_msg_status(self):
assert_equals(self.reconnect_handler._msg_status(self.good_msg), 'ok')
assert_equals(self.reconnect_handler._msg_status(self.bad_msg),
'error')
def test_msg_successful(self):
assert_equals(self.reconnect_handler._msg_successful(self.good_msg),
True)
assert_equals(self.reconnect_handler._msg_successful(self.bad_msg),
False)
def test_msg_error(self):
assert_equals(self.reconnect_handler._msg_error(self.good_msg), None)
assert_equals(self.reconnect_handler._msg_error(self.bad_msg),
u'{}:\n{}'.format('SyntaxError', 'oh no!'))
@gen_test
def test_post_no_json(self):
self.reconnect_handler.request.body = "{{}"
res = yield self.reconnect_handler.post()
assert_equals(res, None)
msg = "Invalid JSON in request body."
self.reconnect_handler.set_status.assert_called_once_with(400)
self.reconnect_handler.finish.assert_called_once_with(msg)
self.spark_events.emit_cluster_change_event.assert_called_once_with(
None, 400, False, msg)
@gen_test
def test_post_no_key(self):
self.reconnect_handler.request.body = json.dumps({})
res = yield self.reconnect_handler.post()
assert_equals(res, None)
msg = 'HTTP 400: Bad Request (Missing argument path)'
self.reconnect_handler.set_status.assert_called_once_with(400)
self.reconnect_handler.finish.assert_called_once_with(msg)
self.spark_events.emit_cluster_change_event.assert_called_once_with(
None, 400, False, msg)
@patch(
'sparkmagic.serverextension.handlers.ReconnectHandler._get_kernel_manager'
)
@gen_test
def test_post_existing_kernel_with_auth_missing_no_auth(
self, _get_kernel_manager):
self.request.body = json.dumps({
"path": self.path,
"username": '',
"password": '',
"endpoint": self.endpoint
})
kernel_manager_future = Future()
kernel_manager_future.set_result(self.individual_kernel_manager)
_get_kernel_manager.return_value = kernel_manager_future
res = yield self.reconnect_handler.post()
assert_equals(res, None)
code = '%{} -s {} -u {} -p {} -t {}'.format(
KernelMagics._do_not_call_change_endpoint.__name__, self.endpoint,
'', '', constants.NO_AUTH)
self.client.execute.assert_called_once_with(code,
silent=False,
store_history=False)
self.reconnect_handler.set_status.assert_called_once_with(200)
self.reconnect_handler.finish.assert_called_once_with(
'{"error": null, "success": true}')
self.spark_events.emit_cluster_change_event.assert_called_once_with(
self.endpoint, 200, True, None)
@patch(
'sparkmagic.serverextension.handlers.ReconnectHandler._get_kernel_manager'
)
@gen_test
def test_post_existing_kernel_with_auth_missing_basic_auth(
self, _get_kernel_manager):
self.request.body = json.dumps({
"path": self.path,
"username": self.username,
"password": self.password,
"endpoint": self.endpoint
})
kernel_manager_future = Future()
kernel_manager_future.set_result(self.individual_kernel_manager)
_get_kernel_manager.return_value = kernel_manager_future
res = yield self.reconnect_handler.post()
assert_equals(res, None)
code = '%{} -s {} -u {} -p {} -t {}'.format(
KernelMagics._do_not_call_change_endpoint.__name__, self.endpoint,
self.username, self.password, constants.AUTH_BASIC)
self.client.execute.assert_called_once_with(code,
silent=False,
store_history=False)
self.reconnect_handler.set_status.assert_called_once_with(200)
self.reconnect_handler.finish.assert_called_once_with(
'{"error": null, "success": true}')
self.spark_events.emit_cluster_change_event.assert_called_once_with(
self.endpoint, 200, True, None)
@patch(
'sparkmagic.serverextension.handlers.ReconnectHandler._get_kernel_manager'
)
@gen_test
def test_post_existing_kernel(self, _get_kernel_manager):
kernel_manager_future = Future()
kernel_manager_future.set_result(self.individual_kernel_manager)
_get_kernel_manager.return_value = kernel_manager_future
res = yield self.reconnect_handler.post()
assert_equals(res, None)
code = '%{} -s {} -u {} -p {} -t {}'.format(
KernelMagics._do_not_call_change_endpoint.__name__, self.endpoint,
self.username, self.password, self.auth)
self.client.execute.assert_called_once_with(code,
silent=False,
store_history=False)
self.reconnect_handler.set_status.assert_called_once_with(200)
self.reconnect_handler.finish.assert_called_once_with(
'{"error": null, "success": true}')
self.spark_events.emit_cluster_change_event.assert_called_once_with(
self.endpoint, 200, True, None)
@patch(
'sparkmagic.serverextension.handlers.ReconnectHandler._get_kernel_manager'
)
@gen_test
def test_post_existing_kernel_failed(self, _get_kernel_manager):
kernel_manager_future = Future()
kernel_manager_future.set_result(self.individual_kernel_manager)
_get_kernel_manager.return_value = kernel_manager_future
self.client.get_shell_msg = MagicMock(return_value=self.bad_msg)
res = yield self.reconnect_handler.post()
assert_equals(res, None)
code = '%{} -s {} -u {} -p {} -t {}'.format(
KernelMagics._do_not_call_change_endpoint.__name__, self.endpoint,
self.username, self.password, self.auth)
self.client.execute.assert_called_once_with(code,
silent=False,
store_history=False)
self.reconnect_handler.set_status.assert_called_once_with(500)
self.reconnect_handler.finish.assert_called_once_with(
'{"error": "SyntaxError:\\noh no!", "success": false}')
self.spark_events.emit_cluster_change_event.assert_called_once_with(
self.endpoint, 500, False, "SyntaxError:\noh no!")
@patch(
'sparkmagic.serverextension.handlers.ReconnectHandler._get_kernel_manager_new_session'
)
@gen_test
def test_get_kernel_manager_no_existing_kernel(
self, _get_kernel_manager_new_session):
different_path = "different_path.ipynb"
km_future = Future()
km_future.set_result(self.individual_kernel_manager)
_get_kernel_manager_new_session.return_value = km_future
km = yield self.reconnect_handler._get_kernel_manager(
different_path, self.kernel_name)
assert_equals(self.individual_kernel_manager, km)
self.individual_kernel_manager.restart_kernel.assert_not_called()
self.kernel_manager.get_kernel.assert_not_called()
_get_kernel_manager_new_session.assert_called_once_with(
different_path, self.kernel_name)
@patch(
'sparkmagic.serverextension.handlers.ReconnectHandler._get_kernel_manager_new_session'
)
@gen_test
def test_get_kernel_manager_existing_kernel(
self, _get_kernel_manager_new_session):
km = yield self.reconnect_handler._get_kernel_manager(
self.path, self.kernel_name)
assert_equals(self.individual_kernel_manager, km)
self.individual_kernel_manager.restart_kernel.assert_called_once_with()
_get_kernel_manager_new_session.assert_not_called()
@patch(
'sparkmagic.serverextension.handlers.ReconnectHandler._get_kernel_manager_new_session'
)
@gen_test
def test_get_kernel_manager_different_kernel_type(
self, _get_kernel_manager_new_session):
different_kernel = "sparkkernel"
km_future = Future()
km_future.set_result(self.individual_kernel_manager)
_get_kernel_manager_new_session.return_value = km_future
km = yield self.reconnect_handler._get_kernel_manager(
self.path, different_kernel)
assert_equals(self.individual_kernel_manager, km)
self.individual_kernel_manager.restart_kernel.assert_not_called()
self.kernel_manager.get_kernel.assert_not_called()
_get_kernel_manager_new_session.assert_called_once_with(
self.path, different_kernel)
self.session_manager.delete_session.assert_called_once_with(
self.session_id)