def test_execute_failure_post_statement_emits_event():
spark_events = MagicMock()
kind = SESSION_KIND_SPARK
http_client = MagicMock()
http_client.get_statement.return_value = tls.TestLivySession.ready_statement_json
conf.override_all({"status_sleep_seconds": 0.01, "statement_sleep_seconds": 0.01})
session = _create_session(kind=kind, http_client=http_client)
session.wait_for_idle = MagicMock()
conf.override_all({})
session.start()
session.wait_for_idle = MagicMock()
command = Command("command", spark_events=spark_events)
http_client.post_statement.side_effect = KeyError("Something bad happened here")
try:
result = command.execute(session)
assert False
except KeyError as e:
spark_events.emit_statement_execution_start_event.assert_called_once_with(
session.guid, session.kind, session.id, command.guid
)
spark_events.emit_statement_execution_end_event._assert_called_once_with(
session.guid, session.kind, session.id, command.guid, -1, False, "KeyError", "Something bad happened here"
)
assert_equals(e, http_client.post_statement.side_effect)
def test_execute_failure_get_statement_output_emits_event():
spark_events = MagicMock()
kind = SESSION_KIND_SPARK
http_client = MagicMock()
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = _create_session(kind=kind, http_client=http_client)
session.wait_for_idle = MagicMock()
conf.override_all({})
session.start(create_sql_context=False)
session.wait_for_idle = MagicMock()
command = Command("command", spark_events=spark_events)
command._get_statement_output = MagicMock(side_effect=AttributeError('OHHHH'))
try:
result = command.execute(session)
assert False
except AttributeError as e:
spark_events.emit_statement_execution_start_event._assert_called_once_with(session.guid, session.kind,
session.id, command.guid)
spark_events.emit_statement_execution_start_event._assert_called_once_with(session.guid, session.kind,
session.id, command.guid,
-1, False, "AttributeError",
"OHHHH")
assert_equals(e, command._get_statement_output.side_effect)
def test_start_emits_start_end_failed_session_when_wait_for_idle_throws(
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
session = self._create_session(kind=kind)
session.wait_for_idle = MagicMock(side_effect=ValueError)
try:
session.start()
assert False
except ValueError:
pass
conf.override_all({})
self.spark_events.emit_session_creation_start_event.assert_called_once_with(
session.guid, kind)
self.spark_events.emit_session_creation_end_event.assert_called_once_with(
session.guid, kind, session.id, session.status, False,
"ValueError", "")
def test_execute():
spark_events = MagicMock()
kind = SESSION_KIND_SPARK
http_client = MagicMock()
http_client.post_session.return_value = tls.TestLivySession.session_create_json
http_client.post_statement.return_value = tls.TestLivySession.post_statement_json
http_client.get_session.return_value = tls.TestLivySession.ready_sessions_json
http_client.get_statement.return_value = tls.TestLivySession.ready_statement_json
conf.override_all({"status_sleep_seconds": 0.01, "statement_sleep_seconds": 0.01})
session = _create_session(kind=kind, http_client=http_client)
conf.override_all({})
session.start()
command = Command("command", spark_events=spark_events)
result = command.execute(session)
http_client.post_statement.assert_called_with(0, {"code": command.code})
http_client.get_statement.assert_called_with(0, 0)
assert result[0]
assert_equals(tls.TestLivySession.pi_result, result[1])
spark_events.emit_statement_execution_start_event.assert_called_once_with(
session.guid, session.kind, session.id, command.guid
)
spark_events.emit_statement_execution_end_event.assert_called_once_with(
session.guid, session.kind, session.id, command.guid, 0, True, "", ""
)
def test_constructor_throws_status_sleep_seconds(self):
conf.override_all({
"status_sleep_seconds": 0,
"statement_sleep_seconds": 2
})
self._create_session()
conf.override_all({})
def test_delete_session_emits_start_failed_end_when_delete_throws(self):
self.http_client.delete_session.side_effect = ValueError
self.http_client.get_session.return_value = self.ready_sessions_json
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = self._create_session()
conf.override_all({})
session.start(create_sql_context=False)
session.id = 0
end_id = session.id
end_status = constants.BUSY_SESSION_STATUS
session.status = end_status
try:
session.delete()
assert False
except ValueError:
pass
self.spark_events.emit_session_deletion_start_event.assert_called_once_with(
session.guid, session.kind, end_id, end_status)
self.spark_events.emit_session_deletion_end_event.assert_called_once_with(
session.guid, session.kind, end_id, end_status, False, "ValueError", "")
def test_constructor_throws_statement_sleep_seconds(self):
conf.override_all({
"status_sleep_seconds": 3,
"statement_sleep_seconds": 0
})
self._create_session()
conf.override_all({})
def test_execute_failure_get_statement_output_emits_event():
spark_events = MagicMock()
kind = SESSION_KIND_SPARK
http_client = MagicMock()
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = _create_session(kind=kind, http_client=http_client)
session.wait_for_idle = MagicMock()
conf.override_all({})
session.start(create_sql_context=False)
session.wait_for_idle = MagicMock()
command = Command("command", spark_events=spark_events)
command._get_statement_output = MagicMock(
side_effect=AttributeError('OHHHH'))
try:
result = command.execute(session)
assert False
except AttributeError as e:
spark_events.emit_statement_execution_start_event._assert_called_once_with(
session.guid, session.kind, session.id, command.guid)
spark_events.emit_statement_execution_start_event._assert_called_once_with(
session.guid, session.kind, session.id, command.guid, -1, False,
"AttributeError", "OHHHH")
assert_equals(e, command._get_statement_output.side_effect)
def test_execute_failure_wait_for_session_emits_event():
spark_events = MagicMock()
kind = SESSION_KIND_SPARK
http_client = MagicMock()
http_client.post_session.return_value = tls.TestLivySession.session_create_json
http_client.post_statement.return_value = tls.TestLivySession.post_statement_json
http_client.get_session.return_value = tls.TestLivySession.ready_sessions_json
http_client.get_statement.return_value = tls.TestLivySession.ready_statement_json
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = _create_session(kind=kind, http_client=http_client)
conf.override_all({})
session.start(create_sql_context=False)
session.wait_for_idle = MagicMock(side_effect=ValueError("yo"))
command = Command("command", spark_events=spark_events)
try:
result = command.execute(session)
assert False
except ValueError as e:
spark_events.emit_statement_execution_start_event._assert_called_once_with(session.guid, session.kind,
session.id, command.guid)
spark_events.emit_statement_execution_start_event._assert_called_once_with(session.guid, session.kind,
session.id, command.guid,
-1, False, "ValueError", "yo")
assert_equals(e, session.wait_for_idle.side_effect)
def test_execute_failure_post_statement_emits_event():
spark_events = MagicMock()
kind = SESSION_KIND_SPARK
http_client = MagicMock()
http_client.get_statement.return_value = tls.TestLivySession.ready_statement_json
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = _create_session(kind=kind, http_client=http_client)
session.wait_for_idle = MagicMock()
conf.override_all({})
session.start()
session.wait_for_idle = MagicMock()
command = Command("command", spark_events=spark_events)
http_client.post_statement.side_effect = KeyError(
'Something bad happened here')
try:
result = command.execute(session)
assert False
except KeyError as e:
spark_events.emit_statement_execution_start_event.assert_called_once_with(
session.guid, session.kind, session.id, command.guid)
spark_events.emit_statement_execution_end_event._assert_called_once_with(
session.guid, session.kind, session.id, command.guid, -1, False,
"KeyError", "Something bad happened here")
assert_equals(e, http_client.post_statement.side_effect)
def test_execute_failure_wait_for_session_emits_event():
spark_events = MagicMock()
kind = SESSION_KIND_SPARK
http_client = MagicMock()
http_client.post_session.return_value = tls.TestLivySession.session_create_json
http_client.post_statement.return_value = tls.TestLivySession.post_statement_json
http_client.get_session.return_value = tls.TestLivySession.ready_sessions_json
http_client.get_statement.return_value = tls.TestLivySession.ready_statement_json
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = _create_session(kind=kind, http_client=http_client)
conf.override_all({})
session.start(create_sql_context=False)
session.wait_for_idle = MagicMock(side_effect=ValueError("yo"))
command = Command("command", spark_events=spark_events)
try:
result = command.execute(session)
assert False
except ValueError as e:
spark_events.emit_statement_execution_start_event._assert_called_once_with(
session.guid, session.kind, session.id, command.guid)
spark_events.emit_statement_execution_start_event._assert_called_once_with(
session.guid, session.kind, session.id, command.guid, -1, False,
"ValueError", "yo")
assert_equals(e, session.wait_for_idle.side_effect)
def test_delete_session_emits_start_failed_end_when_delete_throws(self):
self.http_client.delete_session.side_effect = ValueError
self.http_client.get_session.return_value = self.ready_sessions_json
self.http_client.get_statement.return_value = self.ready_statement_json
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = self._create_session()
conf.override_all({})
session.start()
session.id = 0
end_id = session.id
end_status = constants.BUSY_SESSION_STATUS
session.status = end_status
try:
session.delete()
assert False
except ValueError:
pass
self.spark_events.emit_session_deletion_start_event.assert_called_once_with(
session.guid, session.kind, end_id, end_status)
self.spark_events.emit_session_deletion_end_event.assert_called_once_with(
session.guid, session.kind, end_id, end_status, False, "ValueError", "")
def test_execute():
spark_events = MagicMock()
kind = SESSION_KIND_SPARK
http_client = MagicMock()
http_client.post_session.return_value = tls.TestLivySession.session_create_json
http_client.post_statement.return_value = tls.TestLivySession.post_statement_json
http_client.get_session.return_value = tls.TestLivySession.ready_sessions_json
http_client.get_statement.return_value = tls.TestLivySession.ready_statement_json
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = _create_session(kind=kind, http_client=http_client)
conf.override_all({})
session.start()
command = Command("command", spark_events=spark_events)
result = command.execute(session)
http_client.post_statement.assert_called_with(0, {"code": command.code})
http_client.get_statement.assert_called_with(0, 0)
assert result[0]
assert_equals(tls.TestLivySession.pi_result, result[1])
spark_events.emit_statement_execution_start_event.assert_called_once_with(
session.guid, session.kind, session.id, command.guid)
spark_events.emit_statement_execution_end_event.assert_called_once_with(
session.guid, session.kind, session.id, command.guid, 0, True, "", "")
def test_constructor_throws_invalid_session_sql_combo(self):
conf.override_all({
"status_sleep_seconds": 2,
"statement_sleep_seconds": 2,
"create_sql_context_timeout_seconds": 60
})
self._create_session(sql_created=True)
conf.override_all({})
def test_constructor_throws_invalid_session_sql_combo(self):
conf.override_all({
"status_sleep_seconds": 2,
"statement_sleep_seconds": 2,
"create_sql_context_timeout_seconds": 60
})
self._create_session(sql_created=True)
conf.override_all({})
def test_constructor_throws_status_sleep_seconds(self):
conf.override_all({
"status_sleep_seconds": 0,
"statement_sleep_seconds": 2,
"create_sql_context_timeout_seconds": 60
})
self._create_session()
conf.override_all({})
def test_constructor_throws_status_sleep_seconds(self):
conf.override_all({
"status_sleep_seconds": 0,
"statement_sleep_seconds": 2,
"create_sql_context_timeout_seconds": 60
})
self._create_session()
conf.override_all({})
def _setup():
global magic, spark_controller, shell, ipython_display
conf.override_all({})
magic = RemoteSparkMagics(shell=None, widget=MagicMock())
magic.shell = shell = MagicMock()
magic.ipython_display = ipython_display = MagicMock()
magic.spark_controller = spark_controller = MagicMock()
def test_wrap_unexpected_exceptions_throws_if_all_errors_fatal():
conf.override_all({'all_errors_are_fatal': True})
mock_method = MagicMock(side_effect=ValueError('~~~~~~'))
mock_method.__name__ = 'tos'
decorated = wrap_unexpected_exceptions(mock_method)
assert_equals(decorated.__name__, mock_method.__name__)
result = decorated(self, 'FOOBAR', FOOBAR='FOOBAR')
def test_handle_expected_exceptions_throws_if_all_errors_fatal():
conf.override_all({'all_errors_are_fatal': True})
mock_method = MagicMock(
side_effect=LivyUnexpectedStatusException('Oh no!'))
mock_method.__name__ = 'mock_meth'
decorated = handle_expected_exceptions(mock_method)
assert_equals(decorated.__name__, mock_method.__name__)
result = decorated(self, 1, kwarg='foo')
def test_constructor_starts_with_no_session(self):
conf.override_all({
"status_sleep_seconds": 4,
"statement_sleep_seconds": 2
})
session = self._create_session()
conf.override_all({})
assert session.id == -1
def test_configuration_override_work_with_empty_password():
kpc = { 'username': '******', 'base64_password': '', 'password': '', 'url': '' }
overrides = { conf.kernel_python_credentials.__name__: kpc }
conf.override_all(overrides)
conf.override(conf.status_sleep_seconds.__name__, 1)
assert_equals(conf.d, { conf.kernel_python_credentials.__name__: kpc,
conf.status_sleep_seconds.__name__: 1 })
assert_equals(conf.status_sleep_seconds(), 1)
assert_equals(conf.base64_kernel_python_credentials(), { 'username': '******', 'password': '', 'url': '' })
def test_start_session_times_out_all_errors_are_fatal():
conf.override_all({"all_errors_are_fatal": True})
line = ""
spark_controller.add_session = MagicMock(
side_effect=LivyClientTimeoutException)
assert not magic.session_started
ret = magic._do_not_call_start_session(line)
def test_constructor_starts_with_no_session(self):
conf.override_all({
"status_sleep_seconds": 4,
"statement_sleep_seconds": 2
})
session = self._create_session()
conf.override_all({})
assert session.id == -1
def test_configuration_auth_missing_no_auth():
kpc = {'username': '', 'password': '', 'url': 'L'}
overrides = {conf.kernel_python_credentials.__name__: kpc}
conf.override_all(overrides)
assert_equals(conf.base64_kernel_python_credentials(), {
'username': '',
'password': '',
'url': 'L',
'auth': NO_AUTH
})
def test_custom_headers():
custom_headers = {"header1": "value1"}
overrides = {conf.custom_headers.__name__: custom_headers}
conf.override_all(overrides)
endpoint = Endpoint("http://url.com", constants.NO_AUTH)
client = LivyReliableHttpClient.from_endpoint(endpoint)
headers = client.get_headers()
assert_equals(len(headers), 2)
assert_equals("Content-Type" in headers, True)
assert_equals("header1" in headers, True)
def _setup():
global magic, session, shell, ipython_display
shell = MagicMock()
shell.user_ns = {}
magic = SparkMagicBase(None)
magic.shell = shell
session = MagicMock()
magic.spark_controller = MagicMock()
magic.ipython_display = MagicMock()
conf.override_all({})
def _setup():
global magic, spark_controller, shell, ipython_display, spark_events, conf
conf.override_all({})
spark_events = MagicMock()
magic = TestKernelMagics(shell=None, spark_events=spark_events)
magic.shell = shell = MagicMock()
magic.ipython_display = ipython_display = MagicMock()
magic.spark_controller = spark_controller = MagicMock()
magic._generate_uuid = MagicMock(return_value='0000')
def test_constructor_starts_with_no_session(self):
conf.override_all({
"status_sleep_seconds": 4,
"statement_sleep_seconds": 2,
"create_sql_context_timeout_seconds": 60
})
session = self._create_session()
conf.override_all({})
assert session.id == -1
assert not session.created_sql_context
def test_constructor_starts_heartbeat_with_existing_session(self):
conf.override_all({"heartbeat_refresh_seconds": 0.1})
session_id = 1
session = self._create_session(session_id=session_id)
conf.override_all({})
assert session.id == session_id
assert self.heartbeat_thread.daemon
self.heartbeat_thread.start.assert_called_once_with()
assert not session._heartbeat_thread is None
assert session.properties[constants.LIVY_HEARTBEAT_TIMEOUT_PARAM] > 0
def _setup():
global magic, spark_controller, shell, ipython_display, spark_events, conf
conf.override_all({})
spark_events = MagicMock()
magic = TestKernelMagics(shell=None, spark_events=spark_events)
magic.shell = shell = MagicMock()
magic.ipython_display = ipython_display = MagicMock()
magic.spark_controller = spark_controller = MagicMock()
magic._generate_uuid = MagicMock(return_value='0000')
def test_spark_fatal_spark_statement_exception():
conf.override_all({
"all_errors_are_fatal": True,
})
line = ""
cell = "some spark code"
spark_controller.run_command = MagicMock(side_effect=SparkStatementException('Oh no!'))
magic.spark(line, cell)
def test_constructor_starts_with_no_session(self):
conf.override_all({
"status_sleep_seconds": 4,
"statement_sleep_seconds": 2,
"create_sql_context_timeout_seconds": 60
})
session = self._create_session()
conf.override_all({})
assert session.id == -1
assert not session.created_sql_context
def test_handle_expected_exceptions_handle():
conf.override_all({'all_errors_are_fatal': False})
mock_method = MagicMock(
side_effect=LivyUnexpectedStatusException('ridiculous'))
mock_method.__name__ = 'MockMethod2'
decorated = handle_expected_exceptions(mock_method)
assert_equals(decorated.__name__, mock_method.__name__)
result = decorated(self, 1, kwarg='foo')
assert_is(result, None)
assert_equals(ipython_display.send_error.call_count, 1)
mock_method.assert_called_once_with(self, 1, kwarg='foo')
def test_configuration_override_fallback_to_password():
kpc = {'username': '******', 'password': '******', 'url': 'L'}
overrides = {conf.kernel_python_credentials.__name__: kpc}
conf.override_all(overrides)
conf.override(conf.status_sleep_seconds.__name__, 1)
assert_equals(
conf.d, {
conf.kernel_python_credentials.__name__: kpc,
conf.status_sleep_seconds.__name__: 1
})
assert_equals(conf.status_sleep_seconds(), 1)
assert_equals(conf.base64_kernel_python_credentials(), kpc)
def test_spark_statement_exception():
conf.override_all({
"spark_statement_errors_are_fatal": True
})
mockSparkCommand = MagicMock()
magic._spark_store_command = MagicMock(return_value=mockSparkCommand)
exception = BadUserDataException("Ka-boom!")
magic.spark_controller.run_command.side_effect = [(False, 'out'), exception]
assert_raises(SparkStatementException, magic.execute_spark,"", None, None, None, None, session, True)
magic.spark_controller.cleanup.assert_not_called()
def test_sparkstorecommand_loads_defaults():
defaults = {
conf.default_samplemethod.__name__: "sample",
conf.default_maxrows.__name__: 419,
conf.default_samplefraction.__name__: 0.99,
}
conf.override_all(defaults)
variable_name = "var_name"
sparkcommand = SparkStoreCommand(variable_name)
assert_equals(sparkcommand.samplemethod, defaults[conf.default_samplemethod.__name__])
assert_equals(sparkcommand.maxrows, defaults[conf.default_maxrows.__name__])
assert_equals(sparkcommand.samplefraction, defaults[conf.default_samplefraction.__name__])
def test_delete_session_when_not_started(self):
self.http_client.post_session.return_value = self.session_create_json
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = self._create_session()
conf.override_all({})
session.delete()
assert_equals(session.ipython_display.send_error.call_count, 1)
def test_constructor_starts_with_existing_session(self):
conf.override_all({
"status_sleep_seconds": 4,
"statement_sleep_seconds": 2
})
session_id = 1
session = self._create_session(session_id=session_id, heartbeat_timeout=0)
conf.override_all({})
assert session.id == session_id
assert session._heartbeat_thread is None
assert constants.LIVY_HEARTBEAT_TIMEOUT_PARAM not in list(session.properties.keys())
def test_delete_session_when_not_started(self):
self.http_client.post_session.return_value = self.session_create_json
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = self._create_session()
conf.override_all({})
session.delete()
assert_equals(session.ipython_display.send_error.call_count, 1)
def test_constructor_starts_with_existing_session(self):
conf.override_all({
"status_sleep_seconds": 4,
"statement_sleep_seconds": 2,
"create_sql_context_timeout_seconds": 60
})
session_id = 1
session = self._create_session(session_id=session_id, sql_created=True)
conf.override_all({})
assert session.id == session_id
assert session.created_sql_context
def test_configuration_override_fallback_to_password():
kpc = {'username': '******', 'password': '******', 'url': 'L', 'auth': NO_AUTH}
overrides = {conf.kernel_python_credentials.__name__: kpc}
conf.override_all(overrides)
conf.override(conf.livy_session_startup_timeout_seconds.__name__, 1)
assert_equals(
conf.d, {
conf.kernel_python_credentials.__name__: kpc,
conf.livy_session_startup_timeout_seconds.__name__: 1
})
assert_equals(conf.livy_session_startup_timeout_seconds(), 1)
assert_equals(conf.base64_kernel_python_credentials(), kpc)
def test_sqlquery_loads_defaults():
defaults = {
conf.default_samplemethod.__name__: "sample",
conf.default_maxrows.__name__: 419,
conf.default_samplefraction.__name__: 0.99,
}
conf.override_all(defaults)
query = "DROP TABLE USERS;"
sqlquery = SQLQuery(query)
assert_equals(sqlquery.query, query)
assert_equals(sqlquery.samplemethod, defaults[conf.default_samplemethod.__name__])
assert_equals(sqlquery.maxrows, defaults[conf.default_maxrows.__name__])
assert_equals(sqlquery.samplefraction, defaults[conf.default_samplefraction.__name__])
def test_sqlquery_loads_defaults():
defaults = {
conf.default_samplemethod.__name__: "sample",
conf.default_maxrows.__name__: 419,
conf.default_samplefraction.__name__: 0.99,
}
conf.override_all(defaults)
query = "DROP TABLE USERS;"
sqlquery = SQLQuery(query)
assert_equals(sqlquery.query, query)
assert_equals(sqlquery.samplemethod, defaults[conf.default_samplemethod.__name__])
assert_equals(sqlquery.maxrows, defaults[conf.default_maxrows.__name__])
assert_equals(sqlquery.samplefraction, defaults[conf.default_samplefraction.__name__])
def test_create_sql_hive_context_unknown_throws(self):
kind = "unknown"
self.http_client.post_session.return_value = self.session_create_json
self.http_client.post_statement.return_value = self.post_statement_json
self.http_client.get_session.return_value = self.ready_sessions_json
self.get_statement_responses = [self.running_statement_json, self.ready_statement_json]
self.http_client.get_statement.side_effect = self._next_statement_response_get
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = self._create_session(kind=kind)
conf.override_all({})
session.start()
def test_create_sql_hive_context_unknown_throws(self):
kind = "unknown"
self.http_client.post_session.return_value = self.session_create_json
self.http_client.post_statement.return_value = self.post_statement_json
self.http_client.get_session.return_value = self.ready_sessions_json
self.get_statement_responses = [self.running_statement_json, self.ready_statement_json]
self.http_client.get_statement.side_effect = self._next_statement_response_get
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = self._create_session(kind=kind)
conf.override_all({})
session.start()
def test_constructor_starts_with_existing_session(self):
conf.override_all({
"status_sleep_seconds": 4,
"statement_sleep_seconds": 2,
"create_sql_context_timeout_seconds": 60
})
session_id = 1
session = self._create_session(session_id=session_id, sql_created=True,
should_heartbeat=False)
conf.override_all({})
assert session.id == session_id
assert session.created_sql_context
assert session._heartbeat_thread is None
def test_is_final_status(self):
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = self._create_session()
conf.override_all({})
assert not session.is_final_status("idle")
assert not session.is_final_status("starting")
assert not session.is_final_status("busy")
assert session.is_final_status("dead")
assert session.is_final_status("error")
def test_delete_session_when_active(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
})
session = self._create_session()
conf.override_all({})
session.start(create_sql_context=False)
session.delete()
assert_equals("dead", session.status)
def test_constructor_starts_heartbeat_with_existing_session(self):
conf.override_all({
"status_sleep_seconds": 4,
"statement_sleep_seconds": 2,
"heartbeat_refresh_seconds": 0.1
})
session_id = 1
session = self._create_session(session_id=session_id)
conf.override_all({})
assert session.id == session_id
assert self.heartbeat_thread.daemon
self.heartbeat_thread.start.assert_called_once_with()
assert not session._heartbeat_thread is None
assert session.properties[constants.LIVY_HEARTBEAT_TIMEOUT_PARAM ] > 0
def test_start_with_heartbeat(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
})
session = self._create_session(should_heartbeat=True)
session.create_sql_context = MagicMock()
session.start()
conf.override_all({})
assert self.heartbeat_thread.daemon
self.heartbeat_thread.start.assert_called_once_with()
assert not session._heartbeat_thread is None
def test_logs_gets_latest_logs(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_all_session_logs.return_value = self.log_json
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = self._create_session()
conf.override_all({})
session.start(create_sql_context=False)
logs = session.get_logs()
assert_equals("hi\nhi", logs)
self.http_client.get_all_session_logs.assert_called_with(0)
def test_status_gets_latest_status(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
})
session = self._create_session()
conf.override_all({})
session.start(create_sql_context=False)
session.refresh_status()
state = session.status
assert_equals("idle", state)
self.http_client.get_session.assert_called_with(0)
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_with_heartbeat(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
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01,
})
session = self._create_session()
session.start()
conf.override_all({})
assert self.heartbeat_thread.daemon
self.heartbeat_thread.start.assert_called_once_with()
assert not session._heartbeat_thread is None
assert session.properties[constants.LIVY_HEARTBEAT_TIMEOUT_PARAM ] > 0
def test_start_emits_start_end_session(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
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
kind = constants.SESSION_KIND_SPARK
session = self._create_session(kind=kind)
session.start()
conf.override_all({})
self.spark_events.emit_session_creation_start_event.assert_called_once_with(session.guid, kind)
self.spark_events.emit_session_creation_end_event.assert_called_once_with(
session.guid, kind, session.id, session.status, True, "", "")
def test_constructor_starts_heartbeat_with_existing_session(self):
conf.override_all({
"status_sleep_seconds": 4,
"statement_sleep_seconds": 2,
"create_sql_context_timeout_seconds": 60,
"heartbeat_refresh_seconds": 0.1
})
session_id = 1
session = self._create_session(session_id=session_id, sql_created=True,
should_heartbeat=True)
conf.override_all({})
assert session.id == session_id
assert session.created_sql_context
assert self.heartbeat_thread.daemon
self.heartbeat_thread.start.assert_called_once_with()
assert not session._heartbeat_thread is None
def test_delete_with_heartbeat(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
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = self._create_session()
session.start()
conf.override_all({})
heartbeat_thread = session._heartbeat_thread
session.delete()
self.heartbeat_thread.stop.assert_called_once_with()
assert session._heartbeat_thread is None
def test_start_python_starts_session(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
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
kind = constants.SESSION_KIND_PYSPARK
session = self._create_session(kind=kind)
session.start()
conf.override_all({})
assert_equals(kind, session.kind)
assert_equals("idle", session.status)
assert_equals(0, session.id)
self.http_client.post_session.assert_called_with({"kind": "pyspark", "heartbeatTimeoutInSecond": 60})
def test_create_sql_hive_context_pyspark(self):
kind = constants.SESSION_KIND_PYSPARK
self.http_client.post_session.return_value = self.session_create_json
self.http_client.post_statement.return_value = self.post_statement_json
self.http_client.get_session.return_value = self.ready_sessions_json
self.get_statement_responses = [self.running_statement_json, self.ready_statement_json,
self.ready_statement_json]
self.http_client.get_statement.side_effect = self._next_statement_response_get
conf.override_all({
"status_sleep_seconds": 0.01,
"statement_sleep_seconds": 0.01
})
session = self._create_session(kind=kind)
conf.override_all({})
session.start()
assert call(0, {"code": "from pyspark.sql import HiveContext\n"
"sqlContext = HiveContext(sc)"}) \
in self.http_client.post_statement.call_args_list
