Add more RSpec tests & refactor code