Minor changes and test class