Improve coverage with JUnit tests