Refactor tests
Add test