Minor improvements in test class