Add rspec expectation
Refactor code