Fix broken tests. Finally