Refactor rspec tests and tidy code