Add more specs and some refactoring