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)