Cleaning up query builder and added a couple of extra tests