Improve TestClass