fix imports to work with tests better