and now we have unit tests
refactor using list comprehension and cons