Fixing test failures WIP