def test_poll_for_termination_fail(self, mock_livy):
state_list = 2 * [BatchState.RUNNING] + [BatchState.ERROR]
def side_effect(_):
if state_list:
return state_list.pop(0)
# fail if does not stop right before
raise AssertionError()
mock_livy.side_effect = side_effect
task = LivyOperator(
file='sparkapp',
polling_interval=1,
dag=self.dag,
task_id='livy_example'
)
task._livy_hook = task.get_hook()
with self.assertRaises(AirflowException):
task.poll_for_termination(BATCH_ID)
mock_livy.assert_called_with(BATCH_ID)
self.assertEqual(mock_livy.call_count, 3)
def test_injected_hook(self):
def_hook = LivyHook(livy_conn_id='livyunittest')
task = LivyOperator(file='sparkapp', dag=self.dag, task_id='livy_example')
task._livy_hook = def_hook
assert task.get_hook() == def_hook
def test_execution_with_extra_options(self, mock_post):
extra_options = {'check_response': True}
task = LivyOperator(file='sparkapp',
dag=self.dag,
task_id='livy_example',
extra_options=extra_options)
task.execute(context={})
assert task.get_hook().extra_options == extra_options