Refactor test and update javadoc