def test_get_batch_fail(self, mock):
mock.register_uri('GET',
'//livy:8998/batches/{}'.format(BATCH_ID),
json={'msg': 'Unable to find batch'},
status_code=404,
reason='ERROR')
hook = LivyHook()
with self.assertRaises(AirflowException):
hook.get_batch(BATCH_ID)
def test_get_batch_fail(self, mock):
mock.register_uri(
'GET',
f'//livy:8998/batches/{BATCH_ID}',
json={'msg': 'Unable to find batch'},
status_code=404,
reason='ERROR',
)
hook = LivyHook()
with pytest.raises(AirflowException):
hook.get_batch(BATCH_ID)
def test_get_batch_validation(self, mock):
mock.register_uri('GET',
'//livy:8998/batches/{}'.format(BATCH_ID),
json=SAMPLE_GET_RESPONSE,
status_code=200)
hook = LivyHook()
with self.subTest('get_batch'):
hook.get_batch(BATCH_ID)
# make sure blocked by validation
for val in [None, 'one', {'a': 'b'}]:
with self.subTest('get_batch {}'.format(val)):
with self.assertRaises(TypeError):
hook.get_batch(val)
def test_get_batch_success(self, mock):
mock.register_uri('GET',
'//livy:8998/batches/{}'.format(BATCH_ID),
json={'id': BATCH_ID},
status_code=200)
hook = LivyHook()
resp = hook.get_batch(BATCH_ID)
self.assertIsInstance(resp, dict)
self.assertIn('id', resp)
def test_get_batch_success(self, mock):
mock.register_uri('GET',
f'//livy:8998/batches/{BATCH_ID}',
json={'id': BATCH_ID},
status_code=200)
hook = LivyHook()
resp = hook.get_batch(BATCH_ID)
assert isinstance(resp, dict)
assert 'id' in resp
