def test_run_session_error_during_status_probing(dag, mocker, code): op = LivySessionOperator( statements=[], spill_logs=True, task_id="test_run_session_error_during_status_probing", dag=dag, ) spill_logs_spy = mocker.spy(op, "spill_session_logs") mock_livy_session_responses( mocker, mock_get_session=[ MockedResponse(code, body=f"Response from server:{code}") ], ) with raises(AirflowException) as ae: op.execute({}) print( f"\n\nImitated {code} response from server during session creation probing , " 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() op.spill_logs = False with raises(AirflowException): op.execute({}) # spill_logs=False, but error occured and Operator had the session_id. assert spill_logs_spy.call_count == 2
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_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