Fix broken test suite