def test_link(self):
url = u"https://microsoft.com"
assert_equals(
LivySession.get_html_link(u'Link', url),
u"""<a target="_blank" href="https://microsoft.com">Link</a>""")
url = None
assert_equals(LivySession.get_html_link(u'Link', url), u"")
def test_start_passes_in_all_properties(self):
self.http_client.post_session.return_value = self.session_create_json
self.http_client.get_session.return_value = self.ready_sessions_json
self.http_client.get_statement.return_value = self.ready_statement_json
kind = constants.SESSION_KIND_SPARK
properties = {"kind": kind, "extra": 1}
ipython_display = MagicMock()
session = LivySession(self.http_client, properties, ipython_display)
session.start()
self.http_client.post_session.assert_called_with(properties)
def test_start_passes_in_all_properties(self):
self.http_client.post_session.return_value = self.session_create_json
self.http_client.get_session.return_value = self.ready_sessions_json
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
kind = constants.SESSION_KIND_SPARK
properties = {"kind": kind, "extra": 1}
ipython_display = MagicMock()
session = LivySession(self.http_client, properties, ipython_display)
session.start(create_sql_context=False)
conf.override_all({})
self.http_client.post_session.assert_called_with(properties)
def test_start_passes_in_all_properties(self):
self.http_client.post_session.return_value = self.session_create_json
self.http_client.get_session.return_value = self.ready_sessions_json
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
kind = constants.SESSION_KIND_SPARK
properties = {"kind": kind, "extra": 1}
ipython_display = MagicMock()
session = LivySession(self.http_client, properties, ipython_display)
session.start(create_sql_context=False)
conf.override_all({})
self.http_client.post_session.assert_called_with(properties)
def _create_session(self,
kind=constants.SESSION_KIND_SPARK,
session_id=-1,
heartbeat_timeout=60):
ipython_display = MagicMock()
session = LivySession(self.http_client, {"kind": kind},
ipython_display, session_id, self.spark_events,
heartbeat_timeout, self.heartbeat_thread)
return session
def _create_session(kind=SESSION_KIND_SPARK, session_id=-1,
http_client=None, spark_events=None):
if http_client is None:
http_client = MagicMock()
if spark_events is None:
spark_events = MagicMock()
ipython_display = MagicMock()
session = LivySession(http_client, {"kind": kind, "heartbeatTimeoutInSecond": 60},
ipython_display, session_id, spark_events)
return session
def test_restore_sessions():
with patch('google.auth.default', return_value=(creds, 'project'),\
autospec=True):
spark_controller.get_all_sessions_endpoint = MagicMock(
return_value=sessions_list_mock)
spark_controller.get_managed_clients = []
add_sessions_mock = MagicMock()
spark_controller.session_manager.add_session = add_sessions_mock
add_sessions_mock.assert_called_once_with("my_session", LivySession(http_client=MagicMock(),\
properties={"kind":SESSION_KIND_SPARK, "heartbeatTimeoutInSecond": 60}, ipython_display=ipython_display, session_id=12345))
assert_equals(spark_controller, stored_endpoints)
def _create_session(self, kind=constants.SESSION_KIND_SPARK, session_id=-1, sql_created=False,
should_heartbeat=False):
ipython_display = MagicMock()
session = LivySession(self.http_client,
{"kind": kind},
ipython_display,
session_id,
sql_created,
self.spark_events,
should_heartbeat,
self.heartbeat_thread)
return session
def _create_session(kind=SESSION_KIND_SPARK,
session_id=-1,
sql_created=False,
http_client=None,
spark_events=None):
if http_client is None:
http_client = MagicMock()
if spark_events is None:
spark_events = MagicMock()
ipython_display = MagicMock()
session = LivySession(http_client, {"kind": kind}, ipython_display,
session_id, sql_created, spark_events)
return session
def test_link(self):
url = u"https://microsoft.com"
assert_equals(LivySession.get_html_link(u'Link', url), u"""<a target="_blank" href="https://microsoft.com">Link</a>""")
url = None
assert_equals(LivySession.get_html_link(u'Link', url), u"")
return_value=mocked_db):
global magic, spark_controller, shell, ipython_display
magic = DataprocMagics(shell=None, widget=MagicMock())
magic.shell = shell = MagicMock()
magic.ipython_display = ipython_display = MagicMock()
magic.spark_controller = spark_controller = MagicMock()
def _teardown():
pass
stored_endpoints = ("http://url.com",
Endpoint("http://url.com", "default-credentials"))
get_session_id_to_name = {1234: 'my_session'}
sessions_mock = {'my_session': LivySession(http_client=MagicMock(), properties={"kind":SESSION_KIND_SPARK, \
"heartbeatTimeoutInSecond": 60}, ipython_display=ipython_display, session_id=1234)}
sessions_list_mock = [LivySession(http_client=MagicMock(), properties={"kind":SESSION_KIND_SPARK,\
"heartbeatTimeoutInSecond": 60}, ipython_display=ipython_display, session_id=1234)]
mocked_db = {
'autorestore/stored_endpoints': stored_endpoints,
'autorestore/get_session_id_to_name': get_session_id_to_name,
}
def make_credentials():
return credentials.Credentials(
token=None,
refresh_token='refresh',
token_uri='token_uri',
client_id='client_id',
client_secret='client_secret',
