refactor and fix some broken tests