def test_should_send_single_spark_session_error_message_when_second_sc_starts(
self):
# given
builder = BuilderMock()
spark_session_mock = SingleSparkSessionMock()
engine = SparkEngineMock(builder, spark_session_mock,
SparkSessionFactoryMock())
engine.activate_spark_session()
ipython_manager = IpythonManagerMock()
spark_server_factory = SparkServerFactoryMock()
profile = ProfileMock()
sui = SparkUI2(engine, ipython_manager, spark_server_factory, profile,
CommMock())
msg_start = {
'event': 'start',
'payload': {
"current_profile": "profile1",
"spark_options": {
'spark.executor.memory': '8g',
'spark.master': 'local[10]',
'properties': [{
"name": "wwww",
"value": "wwwww"
}]
}
}
}
# when
sui.handle_msg(sui, msg_start)
# then
self.assertTrue(sui.comm.message["method"] == "update")
error = sui.comm.message["error"]
self.assertTrue(
error["message"] == SparkUI2.ONE_SPARK_SESSION_MSG_ERROR)
def test_should_save_current_profile_when_sc_starts(self):
# given
builder = BuilderMock()
spark_session_mock = SingleSparkSessionMock()
engine = SparkEngineMock(builder, spark_session_mock,
SparkSessionFactoryMock())
ipython_manager = IpythonManagerMock()
spark_server_factory = SparkServerFactoryMock()
profile = ProfileMock()
sui = SparkUI2(engine, ipython_manager, spark_server_factory, profile,
CommMock())
msg_start = {
'event': 'start',
'payload': {
"current_profile": "profile1",
"spark_options": {
'spark.executor.memory': '8g',
'spark.master': 'local[10]',
'properties': []
}
}
}
# when
sui.handle_msg(sui, msg_start)
# then
self.assertTrue(profile.spark_options["current_profile"] == "profile1")
def test_should_save_profiles(self):
# given
builder = BuilderMock()
spark_session_mock = SingleSparkSessionMock()
engine = SparkEngineMock(builder, spark_session_mock,
SparkSessionFactoryMock())
ipython_manager = IpythonManagerMock()
spark_server_factory = SparkServerFactoryMock()
profile = ProfileMock()
sui = SparkUI2(engine, ipython_manager, spark_server_factory, profile,
CommMock())
msg_save_profile = {
"event":
"save_profiles",
"payload": [{
"spark.executor.memory": "8g",
"spark.master": "local[10]",
"name": "new_prof_1",
"spark.executor.cores": "10",
"properties": []
}]
}
# when
sui.handle_msg(sui, msg_save_profile)
# then
result, err = profile.load_profiles()
self.assertTrue(result["profiles"][0]["name"] == "new_prof_1")
self.assertTrue(err is None)
self.assertTrue(sui.comm.message["method"] == "update")
event = sui.comm.message["event"]
self.assertTrue(event["save_profiles"] == "done")
def test_should_not_create_sc_when_builder_is_None(self):
# given
engine = None
spark_server_factory = SparkServerFactoryMock()
ipython = IpythonManagerMock()
profile = ProfileMock()
# when
try:
SparkUI2(engine, ipython, spark_server_factory, profile)
self.fail("builder is None")
except Exception as err:
self.assertTrue("value can not be None" in str(err),
"Should not create SparkUI when builder is None")
def test_should_create_sc(self):
# given
spark_server_factory = SparkServerFactoryMock()
builder = BuilderMock()
spark_session_mock = SingleSparkSessionMock()
engine = SparkEngineMock(builder, spark_session_mock,
SparkSessionFactoryMock())
ipython = IpythonManagerMock()
profile = ProfileMock()
# when
spark_ui = SparkUI2(engine, ipython, spark_server_factory, profile,
CommMock())
# then
self.assertTrue(spark_ui)
def test_should_not_create_sc_when_factory_is_None(self):
# given
builder = BuilderMock()
spark_session_mock = SingleSparkSessionMock()
engine = SparkEngineMock(builder, spark_session_mock,
SparkSessionFactoryMock())
ipython = IpythonManagerMock()
profile = ProfileMock()
spark_server_factory = None
# when
try:
SparkUI2(engine, ipython, spark_server_factory, profile)
self.fail("spark server factory is None")
except Exception as err:
self.assertTrue("value can not be None" in str(err),
"Should not create SparkUI when factory is None")
def test_should_create_spark_conf_based_on_user_conf_when_widget_creation(
self):
# given
builder = BuilderMock()
spark_session_mock = SingleSparkSessionMock()
engine = SparkEngineMock(builder, spark_session_mock,
SparkSessionFactoryMock())
ipython_manager = IpythonManagerMock()
spark_server_factory = SparkServerFactoryMock()
profile = ProfileMock()
# when
sui = SparkUI2(engine, ipython_manager, spark_server_factory, profile,
CommMock())
# then
self.assertTrue(sui.user_spark_conf == {
"name": "",
"prop_1": "user_value_1"
})
def test_should_load_profile_on_widget_creation(self):
# given
builder = BuilderMock()
spark_session_mock = SingleSparkSessionMock()
engine = SparkEngineMock(builder, spark_session_mock,
SparkSessionFactoryMock())
ipython_manager = IpythonManagerMock()
spark_server_factory = SparkServerFactoryMock()
profile = ProfileMock()
# when
sui = SparkUI2(engine, ipython_manager, spark_server_factory, profile,
CommMock())
# then
self.assertTrue(sui.profiles == [{
"name": "",
"prop_1": "init_value_1"
}])
self.assertTrue(sui.current_profile == "")
def test_should_inactivate_single_spark_session_when_application_end(self):
# given
builder = BuilderMock()
spark_session_mock = SingleSparkSessionMock()
engine = SparkEngineMock(builder, spark_session_mock,
SparkSessionFactoryMock())
engine.activate_spark_session()
self.assertTrue(engine.is_active_spark_session())
ipython_manager = IpythonManagerMock()
spark_server_factory = SparkServerFactoryMock()
profile = ProfileMock()
sparkUi = SparkUI2(engine, ipython_manager, spark_server_factory,
profile, CommMock())
listener = SparkListener(sparkUi,
SparkStateProgressUiManagerMock(engine))
# when
listener.onApplicationEnd(None)
# then
self.assertFalse(engine.is_active_spark_session())
def test_should_send_done_message_when_sc_stops(self):
# given
builder = BuilderMock()
spark_session_mock = SingleSparkSessionMock()
engine = SparkEngineMock(builder, spark_session_mock,
SparkSessionFactoryMock())
ipython_manager = IpythonManagerMock()
spark_server_factory = SparkServerFactoryMock()
profile = ProfileMock()
sui = SparkUI2(engine, ipython_manager, spark_server_factory, profile,
CommMock())
msg_start = self.create_msg_start()
sui.handle_msg(sui, msg_start)
msg_stop = {'event': 'stop'}
# when
sui.handle_msg(sui, msg_stop)
# then
self.assertTrue(sui.comm.message["method"] == "update")
event = sui.comm.message["event"]
self.assertTrue(event["stop"] == "done")
def test_auto_connect_spark(self):
# given
ipython_manager = IpythonManagerMock()
builder = BuilderMock()
spark_server_factory = SparkServerFactoryMock()
profile = ProfileMock()
comm_mock = CommMock()
spark_session_mock = SingleSparkSessionMock()
engine = SparkEngineMock(builder, spark_session_mock, SparkSessionFactoryMock())
engine.configure_auto_start()
sui = SparkUI2(engine, ipython_manager, spark_server_factory, profile, comm_mock)
# when
sui.after_display()
# then
self.assertTrue(sui.is_auto_start)
self.assertTrue(sui.comm.message["method"] == "update")
event = sui.comm.message["event"]
self.assertTrue(event["auto_start"] == "done")
self.assertTrue(event["sparkAppId"] == "appIdLocal1")
self.assertTrue(event["sparkUiWebUrl"] == "SparkUiWebUrl1")
def test_should_send_single_spark_session_error_message_when_auto_connect_spark_try_start_second_spark(
self):
# given
ipython_manager = IpythonManagerMock()
builder = BuilderMock()
spark_server_factory = SparkServerFactoryMock()
profile = ProfileMock()
comm_mock = CommMock()
spark_session_mock = SingleSparkSessionMock()
engine = SparkEngineMock(builder, spark_session_mock,
SparkSessionFactoryMock())
engine.activate_spark_session()
engine.configure_auto_start()
sui = SparkUI2(engine, ipython_manager, spark_server_factory, profile,
comm_mock)
# when
sui.after_display()
# then
self.assertTrue(sui.comm.message["method"] == "update")
error = sui.comm.message["error"]
self.assertTrue(
error["message"] == SparkUI2.ONE_SPARK_SESSION_MSG_ERROR)
def _create_spark_ui(self):
spark_widget = SparkUI2(self.spark_engine, self.ipythonManager,
self.server_factory, self.profile, CommMock())
TestSparkUI.spark_widget = spark_widget
return spark_widget