def test_create_sql_context_spark(self): kind = "scala" http_client = MagicMock() self.post_responses = [DummyResponse(201, self.session_create_json), DummyResponse(201, self.post_statement_json), DummyResponse(201, self.post_statement_json)] http_client.post.side_effect = self._next_response_post self.get_responses = [DummyResponse(200, self.ready_sessions_json), DummyResponse(200, self.running_statement_json), DummyResponse(200, self.ready_statement_json), DummyResponse(200, self.ready_sessions_json), DummyResponse(200, self.ready_statement_json)] http_client.get.side_effect = self._next_response_get _t_config_hook({ "status_sleep_seconds": 0.01, "statement_sleep_seconds": 0.01 }) session = LivySession(http_client, kind, "-1", False) _t_config_hook({}) session.start() session.create_sql_context() assert call("/sessions/0/statements", [201], {"code": "val sqlContext = new org.apache.spark.sql.SQLContext" "(sc)\nimport sqlContext.implicits._"}) \ in http_client.post.call_args_list assert call("/sessions/0/statements", [201], {"code": "val hiveContext = new org.apache.spark.sql.hive.Hive" "Context(sc)"}) \ in http_client.post.call_args_list
def test_delete_session_when_dead_throws(self): http_client = MagicMock() http_client.post.return_value = DummyResponse(201, self.session_create_json) session = LivySession(http_client, "scala", state_sleep_seconds=0.01, statement_sleep_seconds=0.01) session._state = "dead" session.delete()
def test_create_sql_hive_context_pyspark(self): kind = "python" http_client = MagicMock() self.post_responses = [DummyResponse(201, self.session_create_json), DummyResponse(201, self.post_statement_json), DummyResponse(201, self.post_statement_json)] http_client.post.side_effect = self._next_response_post self.get_responses = [DummyResponse(200, self.ready_sessions_json), DummyResponse(200, self.running_statement_json), DummyResponse(200, self.ready_statement_json), DummyResponse(200, self.ready_sessions_json), DummyResponse(200, self.ready_statement_json)] http_client.get.side_effect = self._next_response_get _t_config_hook({ "status_sleep_seconds": 0.01, "statement_sleep_seconds": 0.01 }) session = LivySession(http_client, kind, "-1", False) _t_config_hook({}) session.start() session.create_sql_context() assert call("/sessions/0/statements", [201], {"code": "from pyspark.sql import SQLContext\n" "from pyspark.sql.types import *\n" "sqlContext = SQLContext(sc)"}) \ in http_client.post.call_args_list assert call("/sessions/0/statements", [201], {"code": "from pyspark.sql import HiveContext\n" "hiveContext = HiveContext(sc)"}) \ in http_client.post.call_args_list
def test_delete_session_when_not_started(self): http_client = MagicMock() http_client.post.return_value = DummyResponse(201, self.session_create_json) session = LivySession(http_client, "scala", state_sleep_seconds=0.01, statement_sleep_seconds=0.01) session.delete() assert_equals("dead", session._state)
def test_state_gets_latest(self): http_client = MagicMock() http_client.post.return_value = DummyResponse(201, self.session_create_json) http_client.get.return_value = DummyResponse(200, self.ready_sessions_json) session = LivySession(http_client, "scala", state_sleep_seconds=0.01, statement_sleep_seconds=0.01) session.start() state = session.state assert_equals("idle", state) http_client.get.assert_called_with("/sessions", [200])
def test_delete_session_when_dead_throws(self): http_client = MagicMock() http_client.post.return_value = DummyResponse(201, self.session_create_json) _t_config_hook({ "status_sleep_seconds": 0.01, "statement_sleep_seconds": 0.01 }) session = LivySession(http_client, "scala", "-1", False) _t_config_hook({}) session._status = "dead" session.delete()
def test_start_python_starts_session(self): kind = "python" http_client = MagicMock() http_client.post.return_value = DummyResponse(201, self.session_create_json) session = LivySession(http_client, kind, state_sleep_seconds=0.01, statement_sleep_seconds=0.01) session.start() assert_equals(kind, session.language) assert_equals("starting", session._state) assert_equals("0", session._id) http_client.post.assert_called_with("/sessions", [201], {"kind": "pyspark"})
def test_wait_for_state_returns_when_in_state(self): http_client = MagicMock() http_client.post.return_value = DummyResponse(201, self.session_create_json) self.get_responses = [DummyResponse(200, self.busy_sessions_json), DummyResponse(200, self.ready_sessions_json)] http_client.get.side_effect = self._next_response_get session = LivySession(http_client, "scala", state_sleep_seconds=0.01, statement_sleep_seconds=0.01) session.start() session.wait_for_state("idle") http_client.get.assert_called_with("/sessions", [200]) assert_equals(2, http_client.get.call_count)
def test_create_sql_context_unknown_throws(self): kind = "unknown" http_client = MagicMock() self.post_responses = [DummyResponse(201, self.session_create_json), DummyResponse(201, self.post_statement_json)] http_client.post.side_effect = self._next_response_post self.get_responses = [DummyResponse(200, self.ready_sessions_json), DummyResponse(200, self.running_statement_json), DummyResponse(200, self.ready_statement_json)] http_client.get.side_effect = self._next_response_get session = LivySession(http_client, kind, state_sleep_seconds=0.01, statement_sleep_seconds=0.01) session.start() session.create_sql_context()
def test_start_passes_in_all_properties(self): http_client = MagicMock() http_client.post.return_value = DummyResponse(201, self.session_create_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(ipython_display, http_client, "-1", False, properties) session.start() conf.load() http_client.post.assert_called_with("/sessions", [201], properties)
def test_start_passes_in_all_properties(self): self.http_client.post_session.return_value = self.session_create_json self.http_client.get_session.return_value = self.ready_sessions_json conf.override_all({ "status_sleep_seconds": 0.01, "statement_sleep_seconds": 0.01 }) kind = constants.SESSION_KIND_SPARK properties = {"kind": kind, "extra": 1} ipython_display = MagicMock() session = LivySession(self.http_client, properties, ipython_display) session.start(create_sql_context=False) conf.load() self.http_client.post_session.assert_called_with(properties)
def test_start_scala_starts_session(self): kind = "scala" http_client = MagicMock() http_client.post.return_value = DummyResponse(201, self.session_create_json) _t_config_hook({ "status_sleep_seconds": 0.01, "statement_sleep_seconds": 0.01 }) session = LivySession(http_client, kind, "-1", False) session.start() _t_config_hook({}) assert_equals(kind, session.language) assert_equals("starting", session._status) assert_equals("0", session.id) http_client.post.assert_called_with("/sessions", [201], {"kind": "spark"})
def test_start_passes_in_all_properties(self): http_client = MagicMock() http_client.post.return_value = DummyResponse(201, self.session_create_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(http_client, properties, ipython_display) session.start() conf.load() http_client.post.assert_called_with("/sessions", [201], properties)
def test_start_passes_in_all_properties(self): http_client = MagicMock() http_client.post_session.return_value = self.session_create_json 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(http_client, properties, ipython_display) session.start(create_sql_context=False) conf.load() http_client.post_session.assert_called_with(properties)
def test_create_sql_context_spark(self): kind = "scala" http_client = MagicMock() self.post_responses = [DummyResponse(201, self.session_create_json), DummyResponse(201, self.post_statement_json)] http_client.post.side_effect = self._next_response_post self.get_responses = [DummyResponse(200, self.ready_sessions_json), DummyResponse(200, self.running_statement_json), DummyResponse(200, self.ready_statement_json)] http_client.get.side_effect = self._next_response_get session = LivySession(http_client, kind, state_sleep_seconds=0.01, statement_sleep_seconds=0.01) session.start() session.create_sql_context() http_client.post.assert_called_with("/sessions/0/statements", [201], {"code": "val sqlContext = new org.apache.spark.sql.SQLContext(sc)\n" "import sqlContext.implicits._"})
def test_wait_for_status_times_out(self): http_client = MagicMock() http_client.post.return_value = DummyResponse(201, self.session_create_json) self.get_responses = [DummyResponse(200, self.busy_sessions_json), DummyResponse(200, self.busy_sessions_json), DummyResponse(200, self.ready_sessions_json)] http_client.get.side_effect = self._next_response_get _t_config_hook({ "status_sleep_seconds": 0.011, "statement_sleep_seconds": 6000 }) session = LivySession(http_client, "scala", "-1", False) _t_config_hook({}) session.start() session.wait_for_status("idle", 0.01)
def test_execute(self): kind = "scala" http_client = MagicMock() self.post_responses = [DummyResponse(201, self.session_create_json), DummyResponse(201, self.post_statement_json)] http_client.post.side_effect = self._next_response_post self.get_responses = [DummyResponse(200, self.running_statement_json), DummyResponse(200, self.ready_statement_json)] http_client.get.side_effect = self._next_response_get session = LivySession(http_client, kind, state_sleep_seconds=0.01, statement_sleep_seconds=0.01) session.start() command = "command" result = session.execute(command) http_client.post.assert_called_with("/sessions/0/statements", [201], {"code": command}) http_client.get.assert_called_with("/sessions/0/statements", [200]) assert_equals(2, http_client.get.call_count) assert_equals(self.pi_result, result)
def test_create_sql_context_pyspark(self): kind = "python" http_client = MagicMock() self.post_responses = [DummyResponse(201, self.session_create_json), DummyResponse(201, self.post_statement_json)] http_client.post.side_effect = self._next_response_post self.get_responses = [DummyResponse(200, self.ready_sessions_json), DummyResponse(200, self.running_statement_json), DummyResponse(200, self.ready_statement_json)] http_client.get.side_effect = self._next_response_get session = LivySession(http_client, kind, state_sleep_seconds=0.01, statement_sleep_seconds=0.01) session.start() session.create_sql_context() http_client.post.assert_called_with("/sessions/0/statements", [201], {"code": "from pyspark.sql import SQLContext\n" "from pyspark.sql.types import *\n" "sqlContext = SQLContext(sc)"})
def _create_session(self, kind=constants.SESSION_KIND_SPARK, session_id="-1", sql_created=False, http_client=None): if http_client is None: http_client = MagicMock() ipython_display = MagicMock() session = LivySession(http_client, {"kind": kind}, ipython_display, session_id, sql_created) return session
def _create_session(kind=constants.SESSION_KIND_SPARK, session_id=-1, sql_created=False, http_client=None, spark_events=None): if http_client is None: http_client = MagicMock() if spark_events is None: spark_events = MagicMock() ipython_display = MagicMock() session = LivySession(http_client, {"kind": kind}, ipython_display, session_id, sql_created, spark_events) return session
def test_serialize(self): url = "url" username = "******" password = "******" connection_string = get_connection_string(url, username, password) http_client = MagicMock() http_client.connection_string = connection_string kind = "scala" _t_config_hook({ "status_sleep_seconds": 0.01, "statement_sleep_seconds": 0.01 }) session = LivySession(http_client, kind, "-1", False) _t_config_hook({}) serialized = session.get_state().to_dict() assert serialized["connectionstring"] == connection_string assert serialized["id"] == "-1" assert serialized["language"] == kind assert serialized["sqlcontext"] == False assert serialized["version"] == "0.0.0" assert len(serialized.keys()) == 5
def test_is_final_status(self): kind = "scala" http_client = MagicMock() _t_config_hook({ "status_sleep_seconds": 0.01, "statement_sleep_seconds": 0.01 }) session = LivySession(http_client, kind, "-1", False) _t_config_hook({}) 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_create_sql_context_happens_once(self): kind = "scala" http_client = MagicMock() self.post_responses = [DummyResponse(201, self.session_create_json), DummyResponse(201, self.post_statement_json)] http_client.post.side_effect = self._next_response_post self.get_responses = [DummyResponse(200, self.ready_sessions_json), DummyResponse(200, self.running_statement_json), DummyResponse(200, self.ready_statement_json)] http_client.get.side_effect = self._next_response_get session = LivySession(http_client, kind, state_sleep_seconds=0.01, statement_sleep_seconds=0.01) session.start() # Reset the mock so that post called count is accurate http_client.reset_mock() session.create_sql_context() # Second call should not issue a post request session.create_sql_context() http_client.post.assert_called_once_with("/sessions/0/statements", [201], {"code": "val sqlContext = new org.apache.spark.sql.SQLContext(sc)\n" "import sqlContext.implicits._"})