Code clean-up after merge + Rewrote test to use JUnit 4.