def test_run_session_successfully(dag, mocker):
st1 = LivySessionOperator.Statement(kind="spark", code="x = 1;")
st2 = LivySessionOperator.Statement(kind="pyspark", code="print 'hi';")
statements = [st1, st2]
op = LivySessionOperator(
statements=statements,
spill_logs=False,
task_id="test_run_session_successfully",
dag=dag,
)
spill_logs_spy = mocker.spy(op, "spill_session_logs")
submit_statement_spy = mocker.spy(op, "submit_statement")
mock_livy_session_responses(mocker)
op.execute({})
submit_statement_spy.call_count = len(statements)
# spill_logs is False and session completed successfully, so we don't expect logs.
spill_logs_spy.assert_not_called()
op.spill_logs = True
op.execute({})
submit_statement_spy.call_count = len(statements) * 2
# We set spill_logs to True this time, therefore expecting logs.
spill_logs_spy.assert_called_once()
def test_jinja(dag):
st1 = LivySessionOperator.Statement(kind="spark",
code="x=1+{{ custom_param }};")
st2 = LivySessionOperator.Statement(
kind="pyspark", code="print('{{run_id | replace(':', '-')}}')")
op = LivySessionOperator(
name="test_jinja_{{ run_id }}",
statements=[st1, st2],
task_id="test_jinja_session",
dag=dag,
)
op.render_template_fields({"run_id": "hello:world", "custom_param": 3})
assert op.name == "test_jinja_hello:world"
assert op.statements[0].code == "x=1+3;"
assert op.statements[1].code == "print('hello-world')"
def test_allowed_statement_kinds():
kind = "unknown"
with raises(AirflowException) as ae:
LivySessionOperator.Statement(kind=kind, code="a=5;")
print(
f"\n\nTried to create a statement with kind '{kind}', "
f"got the expected exception:\n<{ae.value}>"
)
def test_statement_repr():
st = LivySessionOperator.Statement(kind="spark", code="a=5;\nprint('s');")
assert (st.__str__() == """
{
Statement, kind: spark
code:
--------------------------------------------------------------------------------
a=5;
print('s');
--------------------------------------------------------------------------------
}""")
def test_run_session_error_when_submitting_statement(dag, mocker, code):
st1 = LivySessionOperator.Statement(kind="spark", code="x = 1;")
st2 = LivySessionOperator.Statement(kind="pyspark", code="print 'hi';")
op = LivySessionOperator(
statements=[st1, st2],
spill_logs=True,
task_id="test_run_session_error_when_submitting_statement",
dag=dag,
)
spill_logs_spy = mocker.spy(op, "spill_session_logs")
submit_statement_spy = mocker.spy(op, "submit_statement")
mock_livy_session_responses(
mocker,
mock_post_statement=[
MockedResponse(200, json_body={"id": STATEMENT_ID}),
MockedResponse(code, body=f"Response from server:{code}"),
MockedResponse(200, json_body={"no id here": "haha"}),
],
)
with raises(AirflowException) as ae:
op.execute({})
print(
f"\n\nImitated {code} response from server during second statement submission, "
f"got the expected exception:\n<{ae.value}>")
# spill_logs=True, and Operator had the session_id by the time error occured.
spill_logs_spy.assert_called_once()
assert submit_statement_spy.call_count == 2
op.spill_logs = False
with raises(AirflowException):
op.execute({})
print(
f"\n\nImitated {code} response from server during first statement submission, "
f"got the expected exception:\n<{ae.value}>")
# spill_logs=False, but error occured and Operator had the session_id.
assert spill_logs_spy.call_count == 2
assert submit_statement_spy.call_count == 3
def test_create_session_params(dag, mocker):
heartbeat_timeout = 9
session_start_timeout_sec = 11
session_start_poll_period_sec = 22
statemt_timeout_minutes = 33
statemt_poll_period_sec = 44
http_conn_id = "foo"
spill_logs = True
st1 = LivySessionOperator.Statement(kind="spark", code="x=1")
st2 = LivySessionOperator.Statement(kind="sparkr", code="print hi")
op = LivySessionOperator(
statements=[st1, st2],
kind="pyspark",
proxy_user="******",
jars=["jar1", "jar2"],
py_files=["py_file1", "py_file2"],
files=["file1", "file2"],
driver_memory="driver_memory",
driver_cores=1,
executor_memory="executor_memory",
executor_cores=2,
num_executors=3,
archives=["archive1", "archive2"],
queue="queue",
name="name",
conf={
"key1": "val1",
"key2": 2
},
heartbeat_timeout=heartbeat_timeout,
session_start_timeout_sec=session_start_timeout_sec,
session_start_poll_period_sec=session_start_poll_period_sec,
statemt_timeout_minutes=statemt_timeout_minutes,
statemt_poll_period_sec=statemt_poll_period_sec,
http_conn_id=http_conn_id,
spill_logs=spill_logs,
task_id="test_create_session_params",
dag=dag,
)
mock_response = Response()
mock_response._content = b'{"id": 1}'
patched_hook = mocker.patch.object(HttpHook,
"run",
return_value=mock_response)
op.create_session()
assert op.statements[0] == st1
assert op.statements[1] == st2
assert op.heartbeat_timeout == heartbeat_timeout
assert op.session_start_timeout_sec == session_start_timeout_sec
assert op.session_start_poll_period_sec == session_start_poll_period_sec
assert op.statemt_timeout_minutes == statemt_timeout_minutes
assert op.statemt_poll_period_sec == statemt_poll_period_sec
assert op.http_conn_id == http_conn_id
assert op.spill_logs == spill_logs
expected_json = json.loads("""{
"kind": "pyspark",
"proxyUser": "******",
"jars": [
"jar1",
"jar2"
],
"pyFiles": [
"py_file1",
"py_file2"
],
"files": [
"file1",
"file2"
],
"driverMemory": "driver_memory",
"driverCores": 1,
"executorMemory": "executor_memory",
"executorCores": 2,
"numExecutors": 3,
"archives": [
"archive1",
"archive2"
],
"queue": "queue",
"name": "name",
"conf": {
"key1": "val1",
"key2": 2
},
"heartbeatTimeoutInSecond": 9
}""")
actual_args, actual_kwargs = patched_hook._call_matcher(
patched_hook.call_args)
actual_json = find_json_in_args(actual_args, actual_kwargs)
if actual_json is None:
raise AssertionError(
f"Can not find JSON in HttpHook args.\n"
f"Args:\n{actual_args}\n"
f"KWArgs (JSON should be under 'data' key):\n{actual_kwargs}")
else:
diff = DeepDiff(actual_json, expected_json, ignore_order=True)
if diff:
print(f"\nDifference:\n{json.dumps(diff, indent=2)}")
assert not diff