Fix import of test class in main module.