Refactor and improve test unit