added code to test class