Test your AI Agents with the all-new Agent to Agent Testing Platform.Learn More
How to use _tearDown method in avocado
Best Python code snippet using avocado_python
test_kernel_magics.py
Source:test_kernel_magics.py 
1from mock import MagicMock2from nose.tools import with_setup, raises, assert_equals, assert_is3from IPython.core.magic import magics_class4import sparkmagic.utils.constants as constants5import sparkmagic.utils.configuration as conf6from sparkmagic.kernels.kernelmagics import KernelMagics7from sparkmagic.livyclientlib.exceptions import LivyClientTimeoutException, BadUserDataException,\8 LivyUnexpectedStatusException, SessionManagementException,\9 HttpClientException, DataFrameParseException, SqlContextNotFoundException10from sparkmagic.livyclientlib.endpoint import Endpoint11from sparkmagic.livyclientlib.command import Command12from sparkmagic.utils.constants import NO_AUTH, AUTH_BASIC13magic = None14spark_controller = None15shell = None16ipython_display = MagicMock()17spark_events = None18@magics_class19class TestKernelMagics(KernelMagics):20 def __init__(self, shell, data=None, spark_events=None):21 super(TestKernelMagics, self).__init__(shell, spark_events=spark_events)22 self.language = constants.LANG_PYTHON23 self.endpoint = Endpoint("url", NO_AUTH)24 def refresh_configuration(self):25 self.endpoint = Endpoint("new_url", NO_AUTH)26def _setup():27 global magic, spark_controller, shell, ipython_display, spark_events, conf28 conf.override_all({})29 spark_events = MagicMock()30 magic = TestKernelMagics(shell=None, spark_events=spark_events)31 magic.shell = shell = MagicMock()32 magic.ipython_display = ipython_display = MagicMock()33 magic.spark_controller = spark_controller = MagicMock()34 magic._generate_uuid = MagicMock(return_value='0000')35def _teardown():36 pass37@with_setup(_setup, _teardown)38@raises(NotImplementedError)39def test_local():40 magic.local("")41@with_setup(_setup, _teardown)42def test_start_session():43 line = ""44 assert not magic.session_started45 ret = magic._do_not_call_start_session(line)46 assert ret47 assert magic.session_started48 spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,49 {"kind": constants.SESSION_KIND_PYSPARK})50 # Call a second time51 ret = magic._do_not_call_start_session(line)52 assert ret53 assert magic.session_started54 assert spark_controller.add_session.call_count == 155@with_setup(_setup, _teardown)56def test_start_session_times_out():57 line = ""58 spark_controller.add_session = MagicMock(side_effect=LivyClientTimeoutException)59 assert not magic.session_started60 ret = magic._do_not_call_start_session(line)61 assert not ret62 assert magic.session_started63 assert magic.fatal_error64 assert_equals(ipython_display.send_error.call_count, 1)65 # Call after fatal error66 ipython_display.send_error.reset_mock()67 ret = magic._do_not_call_start_session(line)68 assert not ret69 assert magic.session_started70 assert_equals(ipython_display.send_error.call_count, 1)71@with_setup(_setup, _teardown)72def test_delete_session():73 line = ""74 magic.session_started = True75 magic._do_not_call_delete_session(line)76 assert not magic.session_started77 spark_controller.delete_session_by_name.assert_called_once_with(magic.session_name)78 # Call a second time79 magic._do_not_call_delete_session(line)80 assert not magic.session_started81 assert spark_controller.delete_session_by_name.call_count == 182@with_setup(_setup, _teardown)83def test_delete_session_expected_exception():84 line = ""85 magic.session_started = True86 spark_controller.delete_session_by_name.side_effect = BadUserDataException('hey')87 magic._do_not_call_delete_session(line)88 assert not magic.session_started89 spark_controller.delete_session_by_name.assert_called_once_with(magic.session_name)90 ipython_display.send_error.assert_called_once_with(constants.EXPECTED_ERROR_MSG91 .format(spark_controller.delete_session_by_name.side_effect))92@with_setup(_setup, _teardown)93def test_change_language():94 language = constants.LANG_SCALA.upper()95 line = "-l {}".format(language)96 magic._do_not_call_change_language(line)97 assert_equals(constants.LANG_SCALA, magic.language)98 assert_equals(Endpoint("new_url", NO_AUTH), magic.endpoint)99@with_setup(_setup, _teardown)100def test_change_language_session_started():101 language = constants.LANG_PYTHON102 line = "-l {}".format(language)103 magic.session_started = True104 magic._do_not_call_change_language(line)105 assert_equals(ipython_display.send_error.call_count, 1)106 assert_equals(constants.LANG_PYTHON, magic.language)107 assert_equals(Endpoint("url", NO_AUTH), magic.endpoint)108@with_setup(_setup, _teardown)109def test_change_language_not_valid():110 language = "not_valid"111 line = "-l {}".format(language)112 magic._do_not_call_change_language(line)113 assert_equals(ipython_display.send_error.call_count, 1)114 assert_equals(constants.LANG_PYTHON, magic.language)115 assert_equals(Endpoint("url", NO_AUTH), magic.endpoint)116@with_setup(_setup, _teardown)117def test_change_endpoint():118 u = 'user'119 p = 'password'120 s = 'server'121 t = AUTH_BASIC122 line = "-s {} -u {} -p {} -t {}".format(s, u, p, t)123 magic._do_not_call_change_endpoint(line)124 assert_equals(Endpoint(s, t, u, p), magic.endpoint)125@with_setup(_setup, _teardown)126@raises(BadUserDataException)127def test_change_endpoint_session_started():128 u = 'user'129 p = 'password'130 s = 'server'131 line = "-s {} -u {} -p {}".format(s, u, p)132 magic.session_started = True133 magic._do_not_call_change_endpoint(line)134 135@with_setup(_setup, _teardown)136def test_info():137 magic._print_endpoint_info = print_info_mock = MagicMock()138 line = ""139 session_info = [MagicMock(), MagicMock()]140 spark_controller.get_all_sessions_endpoint = MagicMock(return_value=session_info)141 magic.session_started = True142 magic.info(line)143 print_info_mock.assert_called_once_with(session_info, spark_controller.get_session_id_for_client.return_value)144 spark_controller.get_session_id_for_client.assert_called_once_with(magic.session_name)145 _assert_magic_successful_event_emitted_once('info')146@with_setup(_setup, _teardown)147def test_info_without_active_session():148 magic._print_endpoint_info = print_info_mock = MagicMock()149 line = ""150 session_info = [MagicMock(), MagicMock()]151 spark_controller.get_all_sessions_endpoint = MagicMock(return_value=session_info)152 magic.info(line)153 print_info_mock.assert_called_once_with(session_info, None)154 _assert_magic_successful_event_emitted_once('info')155@with_setup(_setup, _teardown)156def test_info_with_cell_content():157 magic._print_endpoint_info = print_info_mock = MagicMock()158 line = ""159 session_info = ["1", "2"]160 spark_controller.get_all_sessions_endpoint_info = MagicMock(return_value=session_info)161 error_msg = "Cell body for %%info magic must be empty; got 'howdy' instead"162 magic.info(line, cell='howdy')163 print_info_mock.assert_not_called()164 assert_equals(ipython_display.send_error.call_count, 1)165 spark_controller.get_session_id_for_client.assert_not_called()166 _assert_magic_failure_event_emitted_once('info', BadUserDataException(error_msg))167@with_setup(_setup, _teardown)168def test_info_with_argument():169 magic._print_endpoint_info = print_info_mock = MagicMock()170 line = "hey"171 session_info = ["1", "2"]172 spark_controller.get_all_sessions_endpoint_info = MagicMock(return_value=session_info)173 magic.info(line)174 print_info_mock.assert_not_called()175 assert_equals(ipython_display.send_error.call_count, 1)176 spark_controller.get_session_id_for_client.assert_not_called()177@with_setup(_setup, _teardown)178def test_info_unexpected_exception():179 magic._print_endpoint_info = MagicMock()180 line = ""181 spark_controller.get_all_sessions_endpoint = MagicMock(side_effect=ValueError('utter failure'))182 magic.info(line)183 _assert_magic_failure_event_emitted_once('info', spark_controller.get_all_sessions_endpoint.side_effect)184 ipython_display.send_error.assert_called_once_with(constants.INTERNAL_ERROR_MSG185 .format(spark_controller.get_all_sessions_endpoint.side_effect))186@with_setup(_setup, _teardown)187def test_info_expected_exception():188 magic._print_endpoint_info = MagicMock()189 line = ""190 spark_controller.get_all_sessions_endpoint = MagicMock(side_effect=SqlContextNotFoundException('utter failure'))191 magic.info(line)192 _assert_magic_failure_event_emitted_once('info', spark_controller.get_all_sessions_endpoint.side_effect)193 ipython_display.send_error.assert_called_once_with(constants.EXPECTED_ERROR_MSG194 .format(spark_controller.get_all_sessions_endpoint.side_effect))195@with_setup(_setup, _teardown)196def test_help():197 magic.help("")198 assert_equals(ipython_display.html.call_count, 1)199 _assert_magic_successful_event_emitted_once('help')200@with_setup(_setup, _teardown)201def test_help_with_cell_content():202 msg = "Cell body for %%help magic must be empty; got 'HAHAH' instead"203 magic.help("", cell="HAHAH")204 assert_equals(ipython_display.send_error.call_count, 1)205 assert_equals(ipython_display.html.call_count, 0)206 _assert_magic_failure_event_emitted_once('help', BadUserDataException(msg))207@with_setup(_setup, _teardown)208def test_help_with_argument():209 magic.help("argument here")210 assert_equals(ipython_display.send_error.call_count, 1)211 assert_equals(ipython_display.html.call_count, 0)212@with_setup(_setup, _teardown)213def test_logs():214 logs = "logs"215 line = ""216 magic.logs(line)217 ipython_display.write.assert_called_once_with("No logs yet.")218 _assert_magic_successful_event_emitted_once('logs')219 ipython_display.write.reset_mock()220 magic.session_started = True221 spark_controller.get_logs = MagicMock(return_value=logs)222 magic.logs(line)223 ipython_display.write.assert_called_once_with(logs)224 spark_controller.get_logs.assert_called_once_with()225@with_setup(_setup, _teardown)226def test_logs_with_cell_content():227 logs = "logs"228 line = ""229 msg = "Cell body for %%logs magic must be empty; got 'BOOP' instead"230 magic.logs(line, cell="BOOP")231 assert_equals(ipython_display.send_error.call_count, 1)232 _assert_magic_failure_event_emitted_once('logs', BadUserDataException(msg))233@with_setup(_setup, _teardown)234def test_logs_with_argument():235 line = "-h"236 magic.logs(line)237 assert_equals(ipython_display.send_error.call_count, 1)238@with_setup(_setup, _teardown)239def test_logs_unexpected_exception():240 line = ""241 magic.session_started = True242 spark_controller.get_logs = MagicMock(side_effect=SyntaxError('There was some sort of error'))243 magic.logs(line)244 spark_controller.get_logs.assert_called_once_with()245 _assert_magic_failure_event_emitted_once('logs', spark_controller.get_logs.side_effect)246 ipython_display.send_error.assert_called_once_with(constants.INTERNAL_ERROR_MSG247 .format(spark_controller.get_logs.side_effect))248@with_setup(_setup, _teardown)249def test_logs_expected_exception():250 line = ""251 magic.session_started = True252 spark_controller.get_logs = MagicMock(side_effect=LivyUnexpectedStatusException('There was some sort of error'))253 magic.logs(line)254 spark_controller.get_logs.assert_called_once_with()255 _assert_magic_failure_event_emitted_once('logs', spark_controller.get_logs.side_effect)256 ipython_display.send_error.assert_called_once_with(constants.EXPECTED_ERROR_MSG257 .format(spark_controller.get_logs.side_effect))258@with_setup(_setup, _teardown)259def test_configure():260 # Mock info method261 magic.info = MagicMock()262 # Session not started263 conf.override_all({})264 magic.configure('', '{"extra": "yes"}')265 assert_equals(conf.session_configs(), {"extra": "yes"})266 _assert_magic_successful_event_emitted_once('configure')267 magic.info.assert_called_once_with("")268 # Session started - no -f269 magic.session_started = True270 conf.override_all({})271 magic.configure('', "{\"extra\": \"yes\"}")272 assert_equals(conf.session_configs(), {})273 assert_equals(ipython_display.send_error.call_count, 1)274 # Session started - with -f275 magic.info.reset_mock()276 conf.override_all({})277 magic.configure("-f", "{\"extra\": \"yes\"}")278 assert_equals(conf.session_configs(), {"extra": "yes"})279 spark_controller.delete_session_by_name.assert_called_once_with(magic.session_name)280 spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,281 {"kind": constants.SESSION_KIND_PYSPARK, "extra": "yes"})282 magic.info.assert_called_once_with("")283@with_setup(_setup, _teardown)284def test_configure_unexpected_exception():285 magic.info = MagicMock()286 magic._override_session_settings = MagicMock(side_effect=ValueError('help'))287 magic.configure('', '{"extra": "yes"}')288 _assert_magic_failure_event_emitted_once('configure', magic._override_session_settings.side_effect)289 ipython_display.send_error.assert_called_once_with(constants.INTERNAL_ERROR_MSG\290 .format(magic._override_session_settings.side_effect))291@with_setup(_setup, _teardown)292def test_configure_expected_exception():293 magic.info = MagicMock()294 magic._override_session_settings = MagicMock(side_effect=BadUserDataException('help'))295 magic.configure('', '{"extra": "yes"}')296 _assert_magic_failure_event_emitted_once('configure', magic._override_session_settings.side_effect)297 ipython_display.send_error.assert_called_once_with(constants.EXPECTED_ERROR_MSG\298 .format(magic._override_session_settings.side_effect))299@with_setup(_setup, _teardown)300def test_configure_cant_parse_object_as_json():301 magic.info = MagicMock()302 magic._override_session_settings = MagicMock(side_effect=BadUserDataException('help'))303 magic.configure('', "I CAN'T PARSE THIS AS JSON")304 _assert_magic_successful_event_emitted_once('configure')305 assert_equals(ipython_display.send_error.call_count, 1)306@with_setup(_setup, _teardown)307def test_get_session_settings():308 assert magic.get_session_settings("something", False) == "something"309 assert magic.get_session_settings("something ", False) == "something"310 assert magic.get_session_settings(" something", False) == "something"311 assert magic.get_session_settings("-f something", True) == "something"312 assert magic.get_session_settings("something -f", True) == "something"313 assert magic.get_session_settings("something", True) is None314@with_setup(_setup, _teardown)315def test_spark():316 line = ""317 cell = "some spark code"318 spark_controller.run_command = MagicMock(return_value=(True, line))319 magic.spark(line, cell)320 ipython_display.write.assert_called_once_with(line)321 spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,322 {"kind": constants.SESSION_KIND_PYSPARK})323 spark_controller.run_command.assert_called_once_with(Command(cell), None)324@with_setup(_setup, _teardown)325def test_spark_with_argument():326 line = "-z"327 cell = "some spark code"328 spark_controller.run_command = MagicMock(return_value=(True, line))329 magic.spark(line, cell)330 assert_equals(ipython_display.send_error.call_count, 1)331@with_setup(_setup, _teardown)332def test_spark_error():333 line = ""334 cell = "some spark code"335 spark_controller.run_command = MagicMock(return_value=(False, line))336 magic.spark(line, cell)337 ipython_display.send_error.assert_called_once_with(line)338 spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,339 {"kind": constants.SESSION_KIND_PYSPARK})340 spark_controller.run_command.assert_called_once_with(Command(cell), None)341@with_setup(_setup, _teardown)342def test_spark_failed_session_start():343 line = ""344 cell = "some spark code"345 magic._do_not_call_start_session = MagicMock(return_value=False)346 ret = magic.spark(line, cell)347 assert_is(ret, None)348 assert_equals(ipython_display.write.call_count, 0)349 assert_equals(spark_controller.add_session.call_count, 0)350 assert_equals(spark_controller.run_command.call_count, 0)351@with_setup(_setup, _teardown)352def test_spark_unexpected_exception():353 line = ""354 cell = "some spark code"355 spark_controller.run_command = MagicMock(side_effect=Exception('oups'))356 magic.spark(line, cell)357 spark_controller.run_command.assert_called_once_with(Command(cell), None)358 ipython_display.send_error.assert_called_once_with(constants.INTERNAL_ERROR_MSG359 .format(spark_controller.run_command.side_effect))360@with_setup(_setup, _teardown)361def test_spark_expected_exception():362 line = ""363 cell = "some spark code"364 spark_controller.run_command = MagicMock(side_effect=SessionManagementException('oups'))365 magic.spark(line, cell)366 spark_controller.run_command.assert_called_once_with(Command(cell), None)367 ipython_display.send_error.assert_called_once_with(constants.EXPECTED_ERROR_MSG368 .format(spark_controller.run_command.side_effect))369@with_setup(_setup, _teardown)370def test_spark_unexpected_exception_in_storing():371 line = "-o var_name"372 cell = "some spark code"373 side_effect = [(True,'ok'), Exception('oups')]374 spark_controller.run_command = MagicMock(side_effect=side_effect)375 magic.spark(line, cell)376 assert_equals(spark_controller.run_command.call_count, 2)377 spark_controller.run_command.assert_any_call(Command(cell), None)378 ipython_display.send_error.assert_called_with(constants.INTERNAL_ERROR_MSG379 .format(side_effect[1]))380@with_setup(_setup, _teardown)381def test_spark_expected_exception_in_storing():382 line = "-o var_name"383 cell = "some spark code"384 side_effect = [(True,'ok'), SessionManagementException('oups')]385 spark_controller.run_command = MagicMock(side_effect=side_effect)386 magic.spark(line, cell)387 assert spark_controller.run_command.call_count == 2388 spark_controller.run_command.assert_any_call(Command(cell), None)389 ipython_display.send_error.assert_called_with(constants.EXPECTED_ERROR_MSG390 .format(side_effect[1]))391@with_setup(_setup, _teardown)392def test_spark_sample_options():393 line = "-o var_name -m sample -n 142 -r 0.3 -c True"394 cell = ""395 magic.execute_spark = MagicMock()396 ret = magic.spark(line, cell)397 magic.execute_spark.assert_called_once_with(cell, "var_name", "sample", 142, 0.3, None, True)398@with_setup(_setup, _teardown)399def test_spark_false_coerce():400 line = "-o var_name -m sample -n 142 -r 0.3 -c False"401 cell = ""402 magic.execute_spark = MagicMock()403 ret = magic.spark(line, cell)404 magic.execute_spark.assert_called_once_with(cell, "var_name", "sample", 142, 0.3, None, False)405@with_setup(_setup, _teardown)406def test_sql_without_output():407 line = ""408 cell = "some spark code"409 magic.execute_sqlquery = MagicMock()410 magic.sql(line, cell)411 spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,412 {"kind": constants.SESSION_KIND_PYSPARK})413 magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, None, False, None)414@with_setup(_setup, _teardown)415def test_sql_with_output():416 line = "-o my_var"417 cell = "some spark code"418 magic.execute_sqlquery = MagicMock()419 magic.sql(line, cell)420 spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,421 {"kind": constants.SESSION_KIND_PYSPARK})422 magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, "my_var", False, None)423@with_setup(_setup, _teardown)424def test_sql_exception():425 line = "-o my_var"426 cell = "some spark code"427 magic.execute_sqlquery = MagicMock(side_effect=ValueError('HAHAHAHAH'))428 magic.sql(line, cell)429 magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, "my_var", False, None)430 ipython_display.send_error.assert_called_once_with(constants.INTERNAL_ERROR_MSG431 .format(magic.execute_sqlquery.side_effect))432@with_setup(_setup, _teardown)433def test_sql_expected_exception():434 line = "-o my_var"435 cell = "some spark code"436 magic.execute_sqlquery = MagicMock(side_effect=HttpClientException('HAHAHAHAH'))437 magic.sql(line, cell)438 magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, "my_var", False, None)439 ipython_display.send_error.assert_called_once_with(constants.EXPECTED_ERROR_MSG440 .format(magic.execute_sqlquery.side_effect))441@with_setup(_setup, _teardown)442def test_sql_failed_session_start():443 line = ""444 cell = "some spark code"445 magic._do_not_call_start_session = MagicMock(return_value=False)446 ret = magic.sql(line, cell)447 assert_is(ret, None)448 assert_equals(spark_controller.add_session.call_count, 0)449 assert_equals(spark_controller.execute_sqlquery.call_count, 0)450@with_setup(_setup, _teardown)451def test_sql_quiet():452 line = "-q -o Output"453 cell = ""454 magic.execute_sqlquery = MagicMock()455 ret = magic.sql(line, cell)456 spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,457 {"kind": constants.SESSION_KIND_PYSPARK})458 magic.execute_sqlquery.assert_called_once_with(cell, None, None, None, None, "Output", True, None)459@with_setup(_setup, _teardown)460def test_sql_sample_options():461 line = "-q -m sample -n 142 -r 0.3 -c True"462 cell = ""463 magic.execute_sqlquery = MagicMock()464 ret = magic.sql(line, cell)465 spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,466 {"kind": constants.SESSION_KIND_PYSPARK})467 magic.execute_sqlquery.assert_called_once_with(cell, "sample", 142, 0.3, None, None, True, True)468@with_setup(_setup, _teardown)469def test_sql_false_coerce():470 line = "-q -m sample -n 142 -r 0.3 -c False"471 cell = ""472 magic.execute_sqlquery = MagicMock()473 ret = magic.sql(line, cell)474 spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,475 {"kind": constants.SESSION_KIND_PYSPARK})476 magic.execute_sqlquery.assert_called_once_with(cell, "sample", 142, 0.3, None, None, True, False)477@with_setup(_setup, _teardown)478def test_cleanup_without_force():479 line = ""480 cell = ""481 magic.session_started = True482 spark_controller.cleanup_endpoint = MagicMock()483 spark_controller.delete_session_by_name = MagicMock()484 magic.cleanup(line, cell)485 _assert_magic_successful_event_emitted_once('cleanup')486 assert_equals(ipython_display.send_error.call_count, 1)487 assert_equals(spark_controller.cleanup_endpoint.call_count, 0)488@with_setup(_setup, _teardown)489def test_cleanup_with_force():490 line = "-f"491 cell = ""492 magic.session_started = True493 spark_controller.cleanup_endpoint = MagicMock()494 spark_controller.delete_session_by_name = MagicMock()495 magic.cleanup(line, cell)496 _assert_magic_successful_event_emitted_once('cleanup')497 spark_controller.cleanup_endpoint.assert_called_once_with(magic.endpoint)498 spark_controller.delete_session_by_name.assert_called_once_with(magic.session_name)499@with_setup(_setup, _teardown)500def test_cleanup_with_cell_content():501 line = "-f"502 cell = "HEHEHE"503 msg = "Cell body for %%cleanup magic must be empty; got 'HEHEHE' instead"504 magic.session_started = True505 spark_controller.cleanup_endpoint = MagicMock()506 spark_controller.delete_session_by_name = MagicMock()507 magic.cleanup(line, cell)508 assert_equals(ipython_display.send_error.call_count, 1)509 _assert_magic_failure_event_emitted_once('cleanup', BadUserDataException(msg))510@with_setup(_setup, _teardown)511def test_cleanup_exception():512 line = "-f"513 cell = ""514 magic.session_started = True515 spark_controller.cleanup_endpoint = MagicMock(side_effect=ArithmeticError('DIVISION BY ZERO OH NO'))516 magic.cleanup(line, cell)517 _assert_magic_failure_event_emitted_once('cleanup', spark_controller.cleanup_endpoint.side_effect)518 spark_controller.cleanup_endpoint.assert_called_once_with(magic.endpoint)519 ipython_display.send_error.assert_called_once_with(constants.INTERNAL_ERROR_MSG520 .format(spark_controller.cleanup_endpoint.side_effect))521@with_setup(_setup, _teardown)522def test_delete_without_force():523 session_id = 0524 line = "-s {}".format(session_id)525 cell = ""526 spark_controller.delete_session_by_id = MagicMock()527 spark_controller.get_session_id_for_client = MagicMock(return_value=session_id)528 magic.delete(line, cell)529 _assert_magic_successful_event_emitted_once('delete')530 assert_equals(ipython_display.send_error.call_count, 1)531 assert_equals(spark_controller.delete_session_by_id.call_count, 0)532@with_setup(_setup, _teardown)533def test_delete_without_session_id():534 session_id = 0535 line = ""536 cell = ""537 spark_controller.delete_session_by_id = MagicMock()538 spark_controller.get_session_id_for_client = MagicMock(return_value=session_id)539 magic.delete(line, cell)540 _assert_magic_successful_event_emitted_once('delete')541 assert_equals(ipython_display.send_error.call_count, 1)542 assert_equals(spark_controller.delete_session_by_id.call_count, 0)543@with_setup(_setup, _teardown)544def test_delete_with_force_same_session():545 session_id = 0546 line = "-f -s {}".format(session_id)547 cell = ""548 spark_controller.delete_session_by_id = MagicMock()549 spark_controller.get_session_id_for_client = MagicMock(return_value=session_id)550 magic.delete(line, cell)551 _assert_magic_successful_event_emitted_once('delete')552 assert_equals(ipython_display.send_error.call_count, 1)553 assert_equals(spark_controller.delete_session_by_id.call_count, 0)554@with_setup(_setup, _teardown)555def test_delete_with_force_none_session():556 # This happens when session has not been created557 session_id = 0558 line = "-f -s {}".format(session_id)559 cell = ""560 spark_controller.delete_session_by_id = MagicMock()561 spark_controller.get_session_id_for_client = MagicMock(return_value=None)562 magic.delete(line, cell)563 _assert_magic_successful_event_emitted_once('delete')564 spark_controller.get_session_id_for_client.assert_called_once_with(magic.session_name)565 spark_controller.delete_session_by_id.assert_called_once_with(magic.endpoint, session_id)566@with_setup(_setup, _teardown)567def test_delete_with_cell_content():568 # This happens when session has not been created569 session_id = 0570 line = "-f -s {}".format(session_id)571 cell = "~~~"572 msg = "Cell body for %%delete magic must be empty; got '~~~' instead"573 spark_controller.delete_session_by_id = MagicMock()574 spark_controller.get_session_id_for_client = MagicMock(return_value=None)575 magic.delete(line, cell)576 _assert_magic_failure_event_emitted_once('delete', BadUserDataException(msg))577 assert_equals(ipython_display.send_error.call_count, 1)578@with_setup(_setup, _teardown)579def test_delete_with_force_different_session():580 # This happens when session has not been created581 session_id = 0582 line = "-f -s {}".format(session_id)583 cell = ""584 spark_controller.delete_session_by_id = MagicMock()585 spark_controller.get_session_id_for_client = MagicMock(return_value="1")586 magic.delete(line, cell)587 _assert_magic_successful_event_emitted_once('delete')588 spark_controller.get_session_id_for_client.assert_called_once_with(magic.session_name)589 spark_controller.delete_session_by_id.assert_called_once_with(magic.endpoint, session_id)590@with_setup(_setup, _teardown)591def test_delete_exception():592 # This happens when session has not been created593 session_id = 0594 line = "-f -s {}".format(session_id)595 cell = ""596 spark_controller.delete_session_by_id = MagicMock(side_effect=DataFrameParseException('wow'))597 spark_controller.get_session_id_for_client = MagicMock()598 magic.delete(line, cell)599 _assert_magic_failure_event_emitted_once('delete', spark_controller.delete_session_by_id.side_effect)600 spark_controller.get_session_id_for_client.assert_called_once_with(magic.session_name)601 spark_controller.delete_session_by_id.assert_called_once_with(magic.endpoint, session_id)602 ipython_display.send_error.assert_called_once_with(constants.INTERNAL_ERROR_MSG603 .format(spark_controller.delete_session_by_id.side_effect))604@with_setup(_setup, _teardown)605def test_delete_expected_exception():606 # This happens when session has not been created607 session_id = 0608 line = "-f -s {}".format(session_id)609 cell = ""610 spark_controller.delete_session_by_id = MagicMock(side_effect=LivyClientTimeoutException('wow'))611 spark_controller.get_session_id_for_client = MagicMock()612 magic.delete(line, cell)613 _assert_magic_failure_event_emitted_once('delete', spark_controller.delete_session_by_id.side_effect)614 spark_controller.get_session_id_for_client.assert_called_once_with(magic.session_name)615 spark_controller.delete_session_by_id.assert_called_once_with(magic.endpoint, session_id)616 ipython_display.send_error.assert_called_once_with(constants.EXPECTED_ERROR_MSG617 .format(spark_controller.delete_session_by_id.side_effect))618@with_setup(_setup, _teardown)619def test_start_session_displays_fatal_error_when_session_throws():620 e = ValueError("Failed to create the SqlContext.\nError, '{}'".format("Exception"))621 magic.spark_controller.add_session = MagicMock(side_effect=e)622 magic.language = constants.LANG_SCALA623 magic.ipython_display = ipython_display624 magic.ipython_display.send_error = MagicMock()625 magic._do_not_call_start_session("Test Line")626 magic.spark_controller.add_session.assert_called_once_with(magic.session_name, magic.endpoint, False,627 {"kind": constants.SESSION_KIND_SPARK})628 assert magic.fatal_error629 assert magic.fatal_error_message == conf.fatal_error_suggestion().format(str(e))630@with_setup(_setup, _teardown)631def test_kernel_magics_names():632 """The magics machinery in IPython depends on the docstrings and633 method names matching up correctly"""634 assert_equals(magic.help.__name__, 'help')635 assert_equals(magic.local.__name__, 'local')636 assert_equals(magic.info.__name__, 'info')637 assert_equals(magic.logs.__name__, 'logs')638 assert_equals(magic.configure.__name__, 'configure')639 assert_equals(magic.spark.__name__, 'spark')640 assert_equals(magic.sql.__name__, 'sql')641 assert_equals(magic.cleanup.__name__, 'cleanup')642 assert_equals(magic.delete.__name__, 'delete')643def _assert_magic_successful_event_emitted_once(name):644 magic._generate_uuid.assert_called_once_with()645 spark_events.emit_magic_execution_start_event.assert_called_once_with(name, constants.SESSION_KIND_PYSPARK,646 magic._generate_uuid.return_value)647 spark_events.emit_magic_execution_end_event.assert_called_once_with(name, constants.SESSION_KIND_PYSPARK,648 magic._generate_uuid.return_value, True,649 '', '')650def _assert_magic_failure_event_emitted_once(name, error):651 magic._generate_uuid.assert_called_once_with()652 spark_events.emit_magic_execution_start_event.assert_called_once_with(name, constants.SESSION_KIND_PYSPARK,653 magic._generate_uuid.return_value)654 spark_events.emit_magic_execution_end_event.assert_called_once_with(name, constants.SESSION_KIND_PYSPARK,655 magic._generate_uuid.return_value, False,...
test_remotesparkmagics.py
Source:test_remotesparkmagics.py
1from mock import MagicMock2from nose.tools import with_setup3import sparkmagic.utils.configuration as conf4from sparkmagic.utils.constants import EXPECTED_ERROR_MSG, AUTH_BASIC, NO_AUTH5from sparkmagic.magics.remotesparkmagics import RemoteSparkMagics6from sparkmagic.livyclientlib.command import Command7from sparkmagic.livyclientlib.endpoint import Endpoint8from sparkmagic.livyclientlib.exceptions import *9from sparkmagic.livyclientlib.sqlquery import SQLQuery10from sparkmagic.livyclientlib.sparkstorecommand import SparkStoreCommand11magic = None12spark_controller = None13shell = None14ipython_display = None15def _setup():16 global magic, spark_controller, shell, ipython_display17 conf.override_all({})18 magic = RemoteSparkMagics(shell=None, widget=MagicMock())19 magic.shell = shell = MagicMock()20 magic.ipython_display = ipython_display = MagicMock()21 magic.spark_controller = spark_controller = MagicMock()22def _teardown():23 pass24@with_setup(_setup, _teardown)25def test_info_command_parses():26 print_info_mock = MagicMock()27 magic._print_local_info = print_info_mock28 command = "info"29 magic.spark(command)30 print_info_mock.assert_called_once_with()31@with_setup(_setup, _teardown)32def test_info_endpoint_command_parses():33 print_info_mock = MagicMock()34 magic._print_endpoint_info = print_info_mock35 command = "info -u http://microsoft.com"36 spark_controller.get_all_sessions_endpoint_info = MagicMock(return_value=None)37 magic.spark(command)38 print_info_mock.assert_called_once_with(None)39@with_setup(_setup, _teardown)40def test_info_command_exception():41 print_info_mock = MagicMock(side_effect=LivyClientTimeoutException('OHHHHHOHOHOHO'))42 magic._print_local_info = print_info_mock43 command = "info"44 magic.spark(command)45 print_info_mock.assert_called_once_with()46 ipython_display.send_error.assert_called_once_with(EXPECTED_ERROR_MSG47 .format(print_info_mock.side_effect))48@with_setup(_setup, _teardown)49def test_add_sessions_command_parses():50 # Do not skip and python51 add_sessions_mock = MagicMock()52 spark_controller.add_session = add_sessions_mock53 command = "add"54 name = "-s name"55 language = "-l python"56 connection_string = "-u http://url.com -t {} -a sdf -p w".format(AUTH_BASIC)57 line = " ".join([command, name, language, connection_string])58 magic.spark(line)59 add_sessions_mock.assert_called_once_with("name", Endpoint("http://url.com", AUTH_BASIC, "sdf", "w"),60 False, {"kind": "pyspark"})61 # Skip and scala - upper case62 add_sessions_mock = MagicMock()63 spark_controller.add_session = add_sessions_mock64 command = "add"65 name = "-s name"66 language = "-l scala"67 connection_string = "--url http://location:port"68 line = " ".join([command, name, language, connection_string, "-k"])69 magic.spark(line)70 add_sessions_mock.assert_called_once_with("name", Endpoint("http://location:port", NO_AUTH),71 True, {"kind": "spark"})72@with_setup(_setup, _teardown)73def test_add_sessions_command_exception():74 # Do not skip and python75 add_sessions_mock = MagicMock(side_effect=BadUserDataException('hehe'))76 spark_controller.add_session = add_sessions_mock77 command = "add"78 name = "-s name"79 language = "-l python"80 connection_string = "-u http://url.com -t {} -a sdf -p w".format(AUTH_BASIC)81 line = " ".join([command, name, language, connection_string])82 magic.spark(line)83 add_sessions_mock.assert_called_once_with("name", Endpoint("http://url.com", AUTH_BASIC, "sdf", "w"),84 False, {"kind": "pyspark"})85 ipython_display.send_error.assert_called_once_with(EXPECTED_ERROR_MSG86 .format(add_sessions_mock.side_effect))87@with_setup(_setup, _teardown)88def test_add_sessions_command_extra_properties():89 conf.override_all({})90 magic.spark("config", "{\"extra\": \"yes\"}")91 assert conf.session_configs() == {"extra": "yes"}92 add_sessions_mock = MagicMock()93 spark_controller.add_session = add_sessions_mock94 command = "add"95 name = "-s name"96 language = "-l scala"97 connection_string = "-u http://livyendpoint.com"98 line = " ".join([command, name, language, connection_string])99 magic.spark(line)100 add_sessions_mock.assert_called_once_with("name", Endpoint("http://livyendpoint.com", NO_AUTH),101 False, {"kind": "spark", "extra": "yes"})102 conf.override_all({})103@with_setup(_setup, _teardown)104def test_delete_sessions_command_parses():105 mock_method = MagicMock()106 spark_controller.delete_session_by_name = mock_method107 command = "delete -s name"108 magic.spark(command)109 mock_method.assert_called_once_with("name")110 command = "delete -u URL -t {} -a username -p password -i 4".format(AUTH_BASIC)111 mock_method = MagicMock()112 spark_controller.delete_session_by_id = mock_method113 magic.spark(command)114 mock_method.assert_called_once_with(Endpoint("URL", AUTH_BASIC, "username", "password"), 4)115@with_setup(_setup, _teardown)116def test_delete_sessions_command_exception():117 mock_method = MagicMock(side_effect=LivyUnexpectedStatusException('FEEEEEELINGS'))118 spark_controller.delete_session_by_name = mock_method119 command = "delete -s name"120 magic.spark(command)121 mock_method.assert_called_once_with("name")122 ipython_display.send_error.assert_called_once_with(EXPECTED_ERROR_MSG123 .format(mock_method.side_effect))124@with_setup(_setup, _teardown)125def test_cleanup_command_parses():126 mock_method = MagicMock()127 spark_controller.cleanup = mock_method128 line = "cleanup"129 magic.spark(line)130 mock_method.assert_called_once_with()131@with_setup(_setup, _teardown)132def test_cleanup_command_exception():133 mock_method = MagicMock(side_effect=SessionManagementException('Livy did something VERY BAD'))134 spark_controller.cleanup = mock_method135 line = "cleanup"136 magic.spark(line)137 mock_method.assert_called_once_with()138 ipython_display.send_error.assert_called_once_with(EXPECTED_ERROR_MSG139 .format(mock_method.side_effect))140@with_setup(_setup, _teardown)141def test_cleanup_endpoint_command_parses():142 mock_method = MagicMock()143 spark_controller.cleanup_endpoint = mock_method144 line = "cleanup -u endp"145 magic.spark(line)146 mock_method.assert_called_once_with(Endpoint("endp", NO_AUTH))147 line = "cleanup -u endp -a user -p passw -t {}".format(AUTH_BASIC)148 magic.spark(line)149 mock_method.assert_called_with(Endpoint("endp", AUTH_BASIC, "user", "passw"))150@with_setup(_setup, _teardown)151def test_bad_command_writes_error():152 line = "bad_command"153 usage = "Please look at usage of %spark by executing `%spark?`."154 magic.spark(line)155 ipython_display.send_error.assert_called_once_with("Subcommand '{}' not found. {}".format(line, usage))156@with_setup(_setup, _teardown)157def test_run_cell_command_parses():158 run_cell_method = MagicMock()159 result_value = ""160 run_cell_method.return_value = (True, result_value)161 spark_controller.run_command = run_cell_method162 command = "-s"163 name = "sessions_name"164 line = " ".join([command, name])165 cell = "cell code"166 result = magic.spark(line, cell)167 run_cell_method.assert_called_once_with(Command(cell), name)168 assert result is None169 ipython_display.write.assert_called_once_with(result_value)170@with_setup(_setup, _teardown)171def test_run_cell_command_writes_to_err():172 run_cell_method = MagicMock()173 result_value = ""174 run_cell_method.return_value = (False, result_value)175 spark_controller.run_command = run_cell_method176 command = "-s"177 name = "sessions_name"178 line = " ".join([command, name])179 cell = "cell code"180 result = magic.spark(line, cell)181 run_cell_method.assert_called_once_with(Command(cell), name)182 assert result is None183 ipython_display.send_error.assert_called_once_with(result_value)184@with_setup(_setup, _teardown)185def test_run_cell_command_exception():186 run_cell_method = MagicMock()187 run_cell_method.side_effect = HttpClientException('meh')188 spark_controller.run_command = run_cell_method189 command = "-s"190 name = "sessions_name"191 line = " ".join([command, name])192 cell = "cell code"193 result = magic.spark(line, cell)194 run_cell_method.assert_called_once_with(Command(cell), name)195 assert result is None196 ipython_display.send_error.assert_called_once_with(EXPECTED_ERROR_MSG197 .format(run_cell_method.side_effect))198@with_setup(_setup, _teardown)199def test_run_spark_command_parses():200 magic.execute_spark = MagicMock()201 command = "-s"202 name = "sessions_name"203 context = "-c"204 context_name = "spark"205 meth = "-m"206 method_name = "sample"207 line = " ".join([command, name, context, context_name, meth, method_name])208 cell = "cell code"209 result = magic.spark(line, cell)210 magic.execute_spark.assert_called_once_with("cell code",211 None, "sample", None, None, "sessions_name", None)212@with_setup(_setup, _teardown)213def test_run_spark_command_parses_with_coerce():214 magic.execute_spark = MagicMock()215 command = "-s"216 name = "sessions_name"217 context = "-c"218 context_name = "spark"219 meth = "-m"220 method_name = "sample"221 coer = "--coerce"222 coerce_value = "True"223 line = " ".join([command, name, context, context_name, meth, method_name, coer, coerce_value])224 cell = "cell code"225 result = magic.spark(line, cell)226 magic.execute_spark.assert_called_once_with("cell code",227 None, "sample", None, None, "sessions_name", True)228@with_setup(_setup, _teardown)229def test_run_spark_command_parses_with_coerce_false():230 magic.execute_spark = MagicMock()231 command = "-s"232 name = "sessions_name"233 context = "-c"234 context_name = "spark"235 meth = "-m"236 method_name = "sample"237 coer = "--coerce"238 coerce_value = "False"239 line = " ".join([command, name, context, context_name, meth, method_name, coer, coerce_value])240 cell = "cell code"241 result = magic.spark(line, cell)242 magic.execute_spark.assert_called_once_with("cell code",243 None, "sample", None, None, "sessions_name", False)244@with_setup(_setup, _teardown)245def test_run_sql_command_parses_with_coerce_false():246 magic.execute_sqlquery = MagicMock()247 command = "-s"248 name = "sessions_name"249 context = "-c"250 context_name = "sql"251 meth = "-m"252 method_name = "sample"253 coer = "--coerce"254 coerce_value = "False"255 line = " ".join([command, name, context, context_name, meth, method_name, coer, coerce_value])256 cell = "cell code"257 result = magic.spark(line, cell)258 magic.execute_sqlquery.assert_called_once_with("cell code",259 "sample", None, None, "sessions_name", None, False, False)260@with_setup(_setup, _teardown)261def test_run_spark_with_store_command_parses():262 magic.execute_spark = MagicMock()263 command = "-s"264 name = "sessions_name"265 context = "-c"266 context_name = "spark"267 meth = "-m"268 method_name = "sample"269 output = "-o"270 output_var = "var_name"271 line = " ".join([command, name, context, context_name, meth, method_name, output, output_var])272 cell = "cell code"273 result = magic.spark(line, cell)274 magic.execute_spark.assert_called_once_with("cell code",275 "var_name", "sample", None, None, "sessions_name", None)276 277@with_setup(_setup, _teardown)278def test_run_spark_with_store_correct_calls():279 run_cell_method = MagicMock()280 run_cell_method.return_value = (True, "")281 spark_controller.run_command = run_cell_method282 command = "-s"283 name = "sessions_name"284 context = "-c"285 context_name = "spark"286 meth = "-m"287 method_name = "sample"288 output = "-o"289 output_var = "var_name"290 coer = "--coerce"291 coerce_value = "True"292 line = " ".join([command, name, context, context_name, meth, method_name, output, output_var, coer, coerce_value])293 cell = "cell code"294 result = magic.spark(line, cell)295 run_cell_method.assert_any_call(Command(cell), name)296 run_cell_method.assert_any_call(SparkStoreCommand(output_var, samplemethod=method_name, coerce=True), name)297@with_setup(_setup, _teardown)298def test_run_spark_command_exception():299 run_cell_method = MagicMock()300 run_cell_method.side_effect = LivyUnexpectedStatusException('WOW')301 spark_controller.run_command = run_cell_method302 command = "-s"303 name = "sessions_name"304 context = "-c"305 context_name = "spark"306 meth = "-m"307 method_name = "sample"308 output = "-o"309 output_var = "var_name"310 line = " ".join([command, name, context, context_name, meth, method_name, output, output_var])311 cell = "cell code"312 result = magic.spark(line, cell)313 run_cell_method.assert_any_call(Command(cell), name)314 ipython_display.send_error.assert_called_once_with(EXPECTED_ERROR_MSG315 .format(run_cell_method.side_effect))316@with_setup(_setup, _teardown)317def test_run_spark_command_exception_while_storing():318 run_cell_method = MagicMock()319 exception = LivyUnexpectedStatusException('WOW')320 run_cell_method.side_effect = [(True, ""), exception]321 spark_controller.run_command = run_cell_method322 command = "-s"323 name = "sessions_name"324 context = "-c"325 context_name = "spark"326 meth = "-m"327 method_name = "sample"328 output = "-o"329 output_var = "var_name"330 line = " ".join([command, name, context, context_name, meth, method_name, output, output_var])331 cell = "cell code"332 result = magic.spark(line, cell)333 run_cell_method.assert_any_call(Command(cell), name)334 run_cell_method.assert_any_call(SparkStoreCommand(output_var, samplemethod=method_name), name)335 ipython_display.write.assert_called_once_with("")336 ipython_display.send_error.assert_called_once_with(EXPECTED_ERROR_MSG337 .format(exception))338@with_setup(_setup, _teardown)339def test_run_sql_command_parses():340 run_cell_method = MagicMock()341 run_cell_method.return_value = (True, "")342 spark_controller.run_sqlquery = run_cell_method343 command = "-s"344 name = "sessions_name"345 context = "-c"346 context_name = "sql"347 meth = "-m"348 method_name = "sample"349 line = " ".join([command, name, context, context_name, meth, method_name])350 cell = "cell code"351 result = magic.spark(line, cell)352 run_cell_method.assert_called_once_with(SQLQuery(cell, samplemethod=method_name), name)353 assert result is not None354@with_setup(_setup, _teardown)355def test_run_sql_command_exception():356 run_cell_method = MagicMock()357 run_cell_method.side_effect = LivyUnexpectedStatusException('WOW')358 spark_controller.run_sqlquery = run_cell_method359 command = "-s"360 name = "sessions_name"361 context = "-c"362 context_name = "sql"363 meth = "-m"364 method_name = "sample"365 line = " ".join([command, name, context, context_name, meth, method_name])366 cell = "cell code"367 result = magic.spark(line, cell)368 run_cell_method.assert_called_once_with(SQLQuery(cell, samplemethod=method_name), name)369 ipython_display.send_error.assert_called_once_with(EXPECTED_ERROR_MSG370 .format(run_cell_method.side_effect))371@with_setup(_setup, _teardown)372def test_run_sql_command_knows_how_to_be_quiet():373 run_cell_method = MagicMock()374 run_cell_method.return_value = (True, "")375 spark_controller.run_sqlquery = run_cell_method376 command = "-s"377 name = "sessions_name"378 context = "-c"379 context_name = "sql"380 quiet = "-q"381 meth = "-m"382 method_name = "sample"383 line = " ".join([command, name, context, context_name, quiet, meth, method_name])384 cell = "cell code"385 result = magic.spark(line, cell)386 run_cell_method.assert_called_once_with(SQLQuery(cell, samplemethod=method_name), name)387 assert result is None388@with_setup(_setup, _teardown)389def test_logs_subcommand():390 get_logs_method = MagicMock()391 result_value = ""392 get_logs_method.return_value = result_value393 spark_controller.get_logs = get_logs_method394 command = "logs -s"395 name = "sessions_name"396 line = " ".join([command, name])397 cell = "cell code"398 result = magic.spark(line, cell)399 get_logs_method.assert_called_once_with(name)400 assert result is None401 ipython_display.write.assert_called_once_with(result_value)402@with_setup(_setup, _teardown)403def test_logs_exception():404 get_logs_method = MagicMock(side_effect=LivyUnexpectedStatusException('How did this happen?'))405 result_value = ""406 get_logs_method.return_value = result_value407 spark_controller.get_logs = get_logs_method408 command = "logs -s"409 name = "sessions_name"410 line = " ".join([command, name])411 cell = "cell code"412 result = magic.spark(line, cell)413 get_logs_method.assert_called_once_with(name)414 assert result is None415 ipython_display.send_error.assert_called_once_with(EXPECTED_ERROR_MSG...
test_sparkcontroller.py
Source:test_sparkcontroller.py
1from mock import MagicMock, patch2from nose.tools import with_setup, assert_equals, raises3import json4from sparkmagic.livyclientlib.sparkcontroller import SparkController5from sparkmagic.livyclientlib.endpoint import Endpoint6from sparkmagic.livyclientlib.exceptions import SessionManagementException, HttpClientException7import sparkmagic.utils.configuration as conf8from sparkmagic.utils.constants import NO_AUTH9client_manager = None10controller = None11ipython_display = None12class DummyResponse:13 def __init__(self, status_code, json_text):14 self._status_code = status_code15 self._json_text = json_text16 def json(self):17 return json.loads(self._json_text)18 @property19 def status_code(self):20 return self._status_code21def _setup():22 global client_manager, controller, ipython_display23 client_manager = MagicMock()24 ipython_display = MagicMock()25 spark_events = MagicMock()26 controller = SparkController(ipython_display)27 controller.session_manager = client_manager28 controller.spark_events = spark_events29def _teardown():30 pass31@with_setup(_setup, _teardown)32def test_add_session():33 name = "name"34 properties = {"kind": "spark"}35 endpoint = Endpoint("http://location:port", NO_AUTH, "name", "word")36 session = MagicMock()37 controller._livy_session = MagicMock(return_value=session)38 controller._http_client = MagicMock(return_value=MagicMock())39 controller.add_session(name, endpoint, False, properties)40 controller._livy_session.assert_called_once_with(controller._http_client.return_value, properties, ipython_display)41 controller.session_manager.add_session.assert_called_once_with(name, session)42 session.start.assert_called_once()43@with_setup(_setup, _teardown)44def test_add_session_skip():45 name = "name"46 language = "python"47 connection_string = "url=http://location:port;username=name;password=word"48 client = "client"49 session = MagicMock()50 controller._livy_session = MagicMock(return_value=session)51 controller._http_client = MagicMock(return_value=client)52 client_manager.get_sessions_list.return_value = [name]53 controller.add_session(name, language, connection_string, True)54 assert controller._livy_session.create_session.call_count == 055 assert controller._http_client.build_client.call_count == 056 assert client_manager.add_session.call_count == 057 assert session.start.call_count == 058@with_setup(_setup, _teardown)59def test_delete_session():60 name = "name"61 controller.delete_session_by_name(name)62 client_manager.delete_client.assert_called_once_with(name)63@with_setup(_setup, _teardown)64def test_cleanup():65 controller.cleanup()66 client_manager.clean_up_all.assert_called_once_with()67@with_setup(_setup, _teardown)68def test_run_cell():69 default_client = MagicMock()70 chosen_client = MagicMock()71 client_manager.get_any_session = MagicMock(return_value=default_client)72 client_manager.get_session = MagicMock(return_value=chosen_client)73 name = "session_name"74 command = MagicMock()75 controller.run_command(command, name)76 command.execute.assert_called_with(chosen_client)77 controller.run_command(command, None)78 command.execute.assert_called_with(default_client)79@with_setup(_setup, _teardown)80def test_run_sql():81 default_client = MagicMock()82 chosen_client = MagicMock()83 client_manager.get_any_session = MagicMock(return_value=default_client)84 client_manager.get_session = MagicMock(return_value=chosen_client)85 name = "session_name"86 sqlquery = MagicMock()87 controller.run_sqlquery(sqlquery, name)88 sqlquery.execute.assert_called_with(chosen_client)89 controller.run_sqlquery(sqlquery, None)90 sqlquery.execute.assert_called_with(default_client)91@with_setup(_setup, _teardown)92def test_get_client_keys():93 controller.get_client_keys()94 client_manager.get_sessions_list.assert_called_once_with()95@with_setup(_setup, _teardown)96def test_get_all_sessions():97 http_client = MagicMock()98 http_client.get_sessions.return_value = json.loads('{"from":0,"total":2,"sessions":[{"id":0,"state":"idle","kind":'99 '"spark","log":[""]}, {"id":1,"state":"busy","kind":"spark","log"'100 ':[""]}]}')101 controller._http_client = MagicMock(return_value=http_client)102 controller._livy_session = MagicMock()103 sessions = controller.get_all_sessions_endpoint("conn_str")104 assert len(sessions) == 2105@with_setup(_setup, _teardown)106def test_cleanup_endpoint():107 s0 = MagicMock()108 s1 = MagicMock()109 controller.get_all_sessions_endpoint = MagicMock(return_value=[s0, s1])110 controller.cleanup_endpoint("conn_str")111 s0.delete.assert_called_once_with()112 s1.delete.assert_called_once_with()113@with_setup(_setup, _teardown)114def test_delete_session_by_id_existent_non_managed():115 http_client = MagicMock()116 http_client.get_session.return_value = json.loads('{"id":0,"state":"starting","kind":"spark","log":[]}')117 controller._http_client = MagicMock(return_value=http_client)118 session = MagicMock()119 controller._livy_session = MagicMock(return_value=session)120 controller.delete_session_by_id("conn_str", 0)121 controller._livy_session.assert_called_once_with(http_client, {"kind": "spark"}, ipython_display, 0)122 session.delete.assert_called_once_with()123@with_setup(_setup, _teardown)124def test_delete_session_by_id_existent_managed():125 name = "name"126 controller.session_manager.get_session_name_by_id_endpoint = MagicMock(return_value=name)127 controller.session_manager.get_sessions_list = MagicMock(return_value=[name])128 controller.delete_session_by_name = MagicMock()129 controller.delete_session_by_id("conn_str", 0)130 controller.delete_session_by_name.assert_called_once_with(name)131@with_setup(_setup, _teardown)132@raises(HttpClientException)133def test_delete_session_by_id_non_existent():134 http_client = MagicMock()135 http_client.get_session.side_effect = HttpClientException136 controller._http_client = MagicMock(return_value=http_client)137 session = MagicMock()138 controller._livy_session = MagicMock(return_value=session)139 controller.delete_session_by_id("conn_str", 0)140@with_setup(_setup, _teardown)141def test_get_app_id():142 chosen_client = MagicMock()143 controller.get_session_by_name_or_default = MagicMock(return_value=chosen_client)144 result = controller.get_app_id()145 assert_equals(result, chosen_client.get_app_id.return_value)146 chosen_client.get_app_id.assert_called_with()147@with_setup(_setup, _teardown)148def test_get_driver_log():149 chosen_client = MagicMock()150 controller.get_session_by_name_or_default = MagicMock(return_value=chosen_client)151 result = controller.get_driver_log_url()152 assert_equals(result, chosen_client.get_driver_log_url.return_value)153 chosen_client.get_driver_log_url.assert_called_with()154@with_setup(_setup, _teardown)155def test_get_logs():156 chosen_client = MagicMock()157 controller.get_session_by_name_or_default = MagicMock(return_value=chosen_client)158 result = controller.get_logs()159 assert_equals(result, chosen_client.get_logs.return_value)160 chosen_client.get_logs.assert_called_with()161@with_setup(_setup, _teardown)162@raises(SessionManagementException)163def test_get_logs_error():164 chosen_client = MagicMock()165 controller.get_session_by_name_or_default = MagicMock(side_effect=SessionManagementException('THERE WAS A SPOOKY GHOST'))166 result = controller.get_logs()167@with_setup(_setup, _teardown)168def test_get_session_id_for_client():169 assert controller.get_session_id_for_client("name") is not None170 client_manager.get_session_id_for_client.assert_called_once_with("name")171@with_setup(_setup, _teardown)172def test_get_spark_ui_url():173 chosen_client = MagicMock()174 controller.get_session_by_name_or_default = MagicMock(return_value=chosen_client)175 result = controller.get_spark_ui_url()176 assert_equals(result, chosen_client.get_spark_ui_url.return_value)177 chosen_client.get_spark_ui_url.assert_called_with()178@with_setup(_setup, _teardown)179def test_add_session_throws_when_session_start_fails():180 name = "name"181 properties = {"kind": "spark"}182 endpoint = Endpoint("http://location:port", NO_AUTH, "name", "word")183 session = MagicMock()184 controller._livy_session = MagicMock(return_value=session)185 controller._http_client = MagicMock(return_value=MagicMock())186 e = ValueError("Failed to create the SqlContext.\nError, '{}'".format("Exception"))187 session.start = MagicMock(side_effect=e)188 try:189 controller.add_session(name, endpoint, False, properties)190 assert False191 except ValueError as ex:192 assert str(ex) == str(e)...
Automation Testing Tutorials
Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.
LambdaTest Learning Hubs:
- JUnit Tutorial
- TestNG Tutorial
- Webdriver Tutorial
- WebDriverIO Tutorial
- Protractor Tutorial
- Selenium 4 Tutorial
- Jenkins Tutorial
- NUnit Tutorial
- Jest Tutorial
- Playwright Tutorial
- Cypress Tutorial
- PyTest Tutorial
YouTube
You could also refer to video tutorials over LambdaTest YouTube channel to get step by step demonstration from industry experts.
Run avocado automation tests on LambdaTest cloud grid
Perform automation testing on 3000+ real desktop and mobile devices online.
Try LambdaTest Now !!
Get 100 minutes of automation test minutes FREE!!
Was this article helpful?
