improve and fix test