Refactor tests
Add method