refactoring screenshot code
fix tests