Unit test improvement + code improvement