Bug fix
Clean up test class