Minor changes on test class