Refactor tests & some polish