def test_get_batch_state_missing(self, mock):
mock.register_uri('GET',
'//livy:8998/batches/{}/state'.format(BATCH_ID),
json={},
status_code=200)
hook = LivyHook()
with self.assertRaises(AirflowException):
hook.get_batch_state(BATCH_ID)
def test_get_batch_state_missing(self, mock):
mock.register_uri('GET',
f'//livy:8998/batches/{BATCH_ID}/state',
json={},
status_code=200)
hook = LivyHook()
with pytest.raises(AirflowException):
hook.get_batch_state(BATCH_ID)
def test_get_batch_state_fail(self, mock):
mock.register_uri('GET',
f'//livy:8998/batches/{BATCH_ID}/state',
json={},
status_code=400,
reason='ERROR')
hook = LivyHook()
with self.assertRaises(AirflowException):
hook.get_batch_state(BATCH_ID)
def test_get_batch_state_validation(self, mock):
mock.register_uri('GET',
'//livy:8998/batches/{}/state'.format(BATCH_ID),
json=SAMPLE_GET_RESPONSE,
status_code=200)
hook = LivyHook()
with self.subTest('get_batch'):
hook.get_batch_state(BATCH_ID)
for val in [None, 'one', {'a': 'b'}]:
with self.subTest('get_batch {}'.format(val)):
with self.assertRaises(TypeError):
hook.get_batch_state(val)
