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._"})
