fix import issue in tests