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