def test_validate_size_format(self):
with self.subTest('lower 1'):
self.assertTrue(LivyHook._validate_size_format('1m'))
with self.subTest('lower 2'):
self.assertTrue(LivyHook._validate_size_format('1mb'))
with self.subTest('upper 1'):
self.assertTrue(LivyHook._validate_size_format('1G'))
with self.subTest('upper 2'):
self.assertTrue(LivyHook._validate_size_format('1GB'))
with self.subTest('snake 1'):
self.assertTrue(LivyHook._validate_size_format('1Gb'))
with self.subTest('fullmatch'):
with self.assertRaises(ValueError):
self.assertTrue(LivyHook._validate_size_format('1Gb foo'))
with self.subTest('missing size'):
with self.assertRaises(ValueError):
self.assertTrue(LivyHook._validate_size_format('10'))
with self.subTest('numeric'):
with self.assertRaises(ValueError):
LivyHook._validate_size_format(1)
with self.subTest('None'):
# noinspection PyTypeChecker
self.assertTrue(LivyHook._validate_size_format(None))
def test_validate_size_format(self):
with self.subTest('lower 1'):
assert LivyHook._validate_size_format('1m')
with self.subTest('lower 2'):
assert LivyHook._validate_size_format('1mb')
with self.subTest('upper 1'):
assert LivyHook._validate_size_format('1G')
with self.subTest('upper 2'):
assert LivyHook._validate_size_format('1GB')
with self.subTest('snake 1'):
assert LivyHook._validate_size_format('1Gb')
with self.subTest('fullmatch'):
with pytest.raises(ValueError):
assert LivyHook._validate_size_format('1Gb foo')
with self.subTest('missing size'):
with pytest.raises(ValueError):
assert LivyHook._validate_size_format('10')
with self.subTest('numeric'):
with pytest.raises(ValueError):
LivyHook._validate_size_format(1) # noqa
with self.subTest('None'):
assert LivyHook._validate_size_format(None) # noqa
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_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_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_delete_batch_fail(self, mock):
mock.register_uri('DELETE',
'//livy:8998/batches/{}'.format(BATCH_ID),
json={},
status_code=400,
reason='ERROR')
hook = LivyHook()
with self.assertRaises(AirflowException):
hook.delete_batch(BATCH_ID)
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_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_delete_batch_fail(self, mock):
mock.register_uri('DELETE',
f'//livy:8998/batches/{BATCH_ID}',
json={},
status_code=400,
reason='ERROR')
hook = LivyHook()
with pytest.raises(AirflowException):
hook.delete_batch(BATCH_ID)
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
def test_parameters_validation(self):
with self.subTest('not a size'):
with self.assertRaises(ValueError):
LivyHook.build_post_batch_body(file='appname',
executor_memory='xxx')
with self.subTest('list of stringables'):
self.assertEqual(
LivyHook.build_post_batch_body(file='appname',
args=['a', 1, 0.1])['args'],
['a', '1', '0.1'])
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_build_body(self):
with self.subTest('minimal request'):
body = LivyHook.build_post_batch_body(file='appname')
self.assertEqual(body, {'file': 'appname'})
with self.subTest('complex request'):
body = LivyHook.build_post_batch_body(
file='appname',
class_name='org.example.livy',
proxy_user='******',
args=['a', '1'],
jars=['jar1', 'jar2'],
files=['file1', 'file2'],
py_files=['py1', 'py2'],
archives=['arch1', 'arch2'],
queue='queue',
name='name',
conf={'a': 'b'},
driver_cores=2,
driver_memory='1M',
executor_memory='1m',
executor_cores='1', # noqa
num_executors='10',
)
self.assertEqual(
body,
{
'file': 'appname',
'className': 'org.example.livy',
'proxyUser': '******',
'args': ['a', '1'],
'jars': ['jar1', 'jar2'],
'files': ['file1', 'file2'],
'pyFiles': ['py1', 'py2'],
'archives': ['arch1', 'arch2'],
'queue': 'queue',
'name': 'name',
'conf': {
'a': 'b'
},
'driverCores': 2,
'driverMemory': '1M',
'executorMemory': '1m',
'executorCores': '1',
'numExecutors': '10',
},
)
def test_injected_hook(self):
def_hook = LivyHook(livy_conn_id='livyunittest')
task = LivyOperator(file='sparkapp', dag=self.dag, task_id='livy_example')
task._livy_hook = def_hook
assert task.get_hook() == def_hook
def test_build_get_hook(self):
connection_url_mapping = {
# id, expected
'default_port': 'http://host',
'default_protocol': 'http://host',
'port_set': 'http://host:1234',
'schema_set': 'zzz://host',
'dont_override_schema': 'http://host',
}
for conn_id, expected in connection_url_mapping.items():
with self.subTest(conn_id):
hook = LivyHook(livy_conn_id=conn_id)
hook.get_conn()
self.assertEqual(hook.base_url, expected)
def test_delete_batch_success(self, mock):
mock.register_uri('DELETE',
'//livy:8998/batches/{}'.format(BATCH_ID),
json={'msg': 'deleted'},
status_code=200)
resp = LivyHook().delete_batch(BATCH_ID)
self.assertEqual(resp, {'msg': 'deleted'})
def test_delete_batch_success(self, mock):
mock.register_uri('DELETE',
f'//livy:8998/batches/{BATCH_ID}',
json={'msg': 'deleted'},
status_code=200)
resp = LivyHook().delete_batch(BATCH_ID)
assert resp == {'msg': 'deleted'}
def get_hook(self) -> LivyHook:
"""
Get valid hook.
:return: hook
:rtype: LivyHook
"""
if self._livy_hook is None or not isinstance(self._livy_hook, LivyHook):
self._livy_hook = LivyHook(livy_conn_id=self._livy_conn_id)
return self._livy_hook
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_post_batch_success(self, mock):
mock.register_uri('POST',
'//livy:8998/batches',
json={
'id': BATCH_ID,
'state': BatchState.STARTING.value,
'log': []
},
status_code=201)
resp = LivyHook().post_batch(file='sparkapp')
self.assertIsInstance(resp, int)
self.assertEqual(resp, BATCH_ID)
def get_hook(self) -> LivyHook:
"""
Get valid hook.
:return: hook
:rtype: LivyHook
"""
if self._livy_hook is None or not isinstance(self._livy_hook, LivyHook):
self._livy_hook = LivyHook(
livy_conn_id=self._livy_conn_id,
extra_headers=self._extra_headers,
extra_options=self._extra_options,
)
return self._livy_hook
def test_get_batch_state_success(self, mock):
running = BatchState.RUNNING
mock.register_uri('GET',
'//livy:8998/batches/{}/state'.format(BATCH_ID),
json={
'id': BATCH_ID,
'state': running.value
},
status_code=200)
state = LivyHook().get_batch_state(BATCH_ID)
self.assertIsInstance(state, BatchState)
self.assertEqual(state, running)
def test_get_batch_state_success(self, mock):
running = BatchState.RUNNING
mock.register_uri(
'GET',
f'//livy:8998/batches/{BATCH_ID}/state',
json={
'id': BATCH_ID,
'state': running.value
},
status_code=200,
)
state = LivyHook().get_batch_state(BATCH_ID)
assert isinstance(state, BatchState)
assert state == running
def test_check_session_id(self):
with self.subTest('valid 00'):
try:
LivyHook._validate_session_id(100)
except TypeError:
self.fail("")
with self.subTest('valid 01'):
try:
LivyHook._validate_session_id(0)
except TypeError:
self.fail("")
with self.subTest('None'):
with self.assertRaises(TypeError):
LivyHook._validate_session_id(None)
with self.subTest('random string'):
with self.assertRaises(TypeError):
LivyHook._validate_session_id('asd')
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)
def test_delete_batch_validation(self, mock):
mock.register_uri('DELETE',
'//livy:8998/batches/{}'.format(BATCH_ID),
json={'id': BATCH_ID},
status_code=200)
hook = LivyHook()
with self.subTest('get_batch'):
hook.delete_batch(BATCH_ID)
for val in [None, 'one', {'a': 'b'}]:
with self.subTest('get_batch {}'.format(val)):
with self.assertRaises(TypeError):
hook.delete_batch(val)
def test_get_batch_validation(self, mock):
mock.register_uri('GET',
f'//livy:8998/batches/{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(f'get_batch {val}'):
with self.assertRaises(TypeError):
hook.get_batch(val)
def test_missing_host(self):
with self.assertRaises(AirflowException):
LivyHook(livy_conn_id='missing_host').get_conn()
