Some import clean up and fixed failing JUnit test