More refactoring and adding some rspec tests