def test_run_batch_successfully(dag, mocker):
op = LivyBatchOperator(spill_logs=False,
task_id="test_run_batch_successfully",
dag=dag)
spill_logs_spy = mocker.spy(op, "spill_batch_logs")
submit_batch_spy = mocker.spy(op, "submit_batch")
mock_livy_batch_responses(mocker)
op.execute({})
submit_batch_spy.assert_called_once()
# spill_logs is False and batch completed successfully, so we don't expect logs.
spill_logs_spy.assert_not_called()
op.spill_logs = True
op.execute({})
# We set spill_logs to True this time, therefore expecting logs.
spill_logs_spy.assert_called_once()
def test_run_batch_error_during_status_probing(dag, mocker, code):
op = LivyBatchOperator(
spill_logs=True,
task_id="test_run_batch_error_during_status_probing",
dag=dag,
)
spill_logs_spy = mocker.spy(op, "spill_batch_logs")
mock_livy_batch_responses(
mocker,
mock_get=[MockedResponse(code, body=f"Response from server:{code}")])
with raises(AirflowException) as ae:
op.execute({})
print(
f"\n\nImitated {code} response from server during batch status probing , "
f"got the expected exception:\n<{ae.value}>")
# spill_logs=True, and Operator had the batch_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 batch_id.
assert spill_logs_spy.call_count == 2