modified to fix test