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