Moved code around and added another test