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')