def test_missing_batch_id(self, mock):
mock.register_uri('POST',
'//livy:8998/batches',
json={},
status_code=201)
hook = LivyHook()
with self.assertRaises(AirflowException):
hook.post_batch(file='sparkapp')
def test_post_batch_fail(self, mock):
mock.register_uri('POST',
'//livy:8998/batches',
json={},
status_code=400,
reason='ERROR')
hook = LivyHook()
with self.assertRaises(AirflowException):
hook.post_batch(file='sparkapp')
def test_post_batch_arguments(self, mock_request):
mock_request.return_value.status_code = 201
mock_request.return_value.json.return_value = {
'id': BATCH_ID,
'state': BatchState.STARTING.value,
'log': []
}
hook = LivyHook()
resp = hook.post_batch(file='sparkapp')
mock_request.assert_called_once_with(method='POST',
endpoint='/batches',
data=json.dumps(
{'file': 'sparkapp'}))
request_args = mock_request.call_args[1]
self.assertIn('data', request_args)
self.assertIsInstance(request_args['data'], str)
self.assertIsInstance(resp, int)
self.assertEqual(resp, BATCH_ID)
def test_invalid_uri(self):
hook = LivyHook(livy_conn_id='invalid_uri')
with self.assertRaises(RequestException):
hook.post_batch(file='sparkapp')
