Apply @andrewsmedina's suggestion and improve tests

@ricardodani you may still squash these two commits if you'd like